64 lines
2.0 KiB
Python
64 lines
2.0 KiB
Python
# 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")
|