feat: add yarn client (status, logs, kill)

tests/unit/test_yarn_client.py

@@ -0,0 +1,63 @@
+# coding=utf-8
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from spark_executor.core.yarn_client import (
+    YarnError,
+    get_application_logs,
+    get_application_status,
+    kill_application,
+)
+
+
+def _fake_proc(returncode: int, stdout: str = "", stderr: str = ""):
+    p = MagicMock()
+    p.returncode = returncode
+    p.stdout = stdout
+    p.stderr = stderr
+    return p
+
+
+def test_status_parses_state_line():
+    fake = _fake_proc(
+        0,
+        stdout="Application Report :\n  State : RUNNING\n  ...\n",
+    )
+    with patch("spark_executor.core.yarn_client.subprocess.run", return_value=fake):
+        state, raw = get_application_status("application_1")
+    assert state == "RUNNING"
+    assert "RUNNING" in raw
+
+
+def test_status_raises_on_nonzero_return():
+    fake = _fake_proc(1, stderr="not found")
+    with patch("spark_executor.core.yarn_client.subprocess.run", return_value=fake):
+        with pytest.raises(YarnError):
+            get_application_status("application_x")
+
+
+def test_logs_returns_stdout():
+    fake = _fake_proc(0, stdout="log line 1\nlog line 2\n")
+    with patch("spark_executor.core.yarn_client.subprocess.run", return_value=fake) as m:
+        out = get_application_logs("application_1")
+    assert out == "log line 1\nlog line 2\n"
+    args = m.call_args.args[0]
+    assert args[:3] == ["yarn", "logs", "-applicationId"]
+    assert args[3] == "application_1"
+
+
+def test_kill_invokes_yarn_application_kill():
+    fake = _fake_proc(0)
+    with patch("spark_executor.core.yarn_client.subprocess.run", return_value=fake) as m:
+        kill_application("application_1")
+    args = m.call_args.args[0]
+    assert args[:3] == ["yarn", "application", "-kill"]
+    assert args[3] == "application_1"
+
+
+def test_kill_raises_on_nonzero_return():
+    fake = _fake_proc(1, stderr="denied")
+    with patch("spark_executor.core.yarn_client.subprocess.run", return_value=fake):
+        with pytest.raises(YarnError):
+            kill_application("application_1")