# coding=utf-8 from unittest.mock import patch import httpx import pytest from common import config from spark_executor.core.yarn_client import ( YarnConfigError, YarnError, get_application_logs, get_application_status, kill_application, ) RM = "http://rm:8088" @pytest.fixture(autouse=True) def _restore_settings(): """Each test gets a clean copy of `settings` so env-var-style overrides in one test don't leak into the next.""" snapshot = config.Settings( data_dir=config.settings.data_dir, jobs_dir=config.settings.jobs_dir, yarn_resource_manager_url=config.settings.yarn_resource_manager_url, log_level=config.settings.log_level, ) yield config.settings.data_dir = snapshot.data_dir config.settings.jobs_dir = snapshot.jobs_dir config.settings.yarn_resource_manager_url = snapshot.yarn_resource_manager_url config.settings.log_level = snapshot.log_level def _resp(status: int, *, json_data=None, text: str | None = None) -> httpx.Response: if json_data is not None: return httpx.Response(status, json=json_data) return httpx.Response(status, text=text or "") # --- get_application_status --- def test_status_parses_app_state(): fake = _resp(200, json_data={"app": {"id": "application_1", "state": "RUNNING"}}) with patch("spark_executor.core.yarn_client.httpx.request", return_value=fake) as m: state, raw = get_application_status("application_1", RM) assert state == "RUNNING" assert "RUNNING" in raw args = m.call_args.args assert args == ("GET", f"{RM}/ws/v1/cluster/apps/application_1") def test_status_raises_on_404(): with patch("spark_executor.core.yarn_client.httpx.request", return_value=_resp(404)): with pytest.raises(YarnError, match="not found"): get_application_status("application_x", RM) def test_status_raises_on_5xx(): with patch( "spark_executor.core.yarn_client.httpx.request", return_value=_resp(503, text="upstream down"), ): with pytest.raises(YarnError, match="503"): get_application_status("application_1", RM) def test_status_raises_when_state_field_missing(): fake = _resp(200, json_data={"app": {"id": "application_1"}}) with patch("spark_executor.core.yarn_client.httpx.request", return_value=fake): with pytest.raises(YarnError, match="Could not parse YARN state"): get_application_status("application_1", RM) def test_status_requires_rm_url(): config.settings.yarn_resource_manager_url = None with pytest.raises(YarnConfigError): get_application_status("application_1", None) def test_status_falls_back_to_settings_yarn_rm_url(): config.settings.yarn_resource_manager_url = "http://env-rm:8088" fake = _resp(200, json_data={"app": {"state": "FINISHED"}}) with patch("spark_executor.core.yarn_client.httpx.request", return_value=fake) as m: state, _ = get_application_status("application_1", None) assert state == "FINISHED" assert "env-rm:8088" in m.call_args.args[1] def test_status_rejects_non_http_url(): with pytest.raises(YarnConfigError, match="must start with"): get_application_status("application_1", "rm:8088") # --- get_application_logs --- def test_logs_returns_text_on_h3_aggregated_endpoint(): fake = _resp(200, text="log line 1\nlog line 2\n") with patch("spark_executor.core.yarn_client.httpx.request", return_value=fake) as m: out = get_application_logs("application_1", RM) assert out == "log line 1\nlog line 2\n" assert m.call_args.args == ( "GET", f"{RM}/ws/v1/cluster/apps/application_1/aggregated-logs", ) def test_logs_falls_back_to_nodemanager_on_404(): responses = [ _resp(404), # aggregated-logs H3 endpoint missing _resp(200, json_data={"appAttempts": {"appAttempt": [{"id": "attempt_1"}]}}), _resp( 200, json_data={ "containers": { "container": [ { "id": "container_1", "user": "hdfs", "nodeHttpAddress": "nm1:8042", } ] } }, ), _resp(200, text="container log content"), ] with patch( "spark_executor.core.yarn_client.httpx.request", side_effect=responses ) as m: out = get_application_logs("application_1", RM) assert "container_1" in out assert "container log content" in out calls = m.call_args_list assert calls[0].args == ( "GET", f"{RM}/ws/v1/cluster/apps/application_1/aggregated-logs", ) assert calls[1].args == ("GET", f"{RM}/ws/v1/cluster/apps/application_1/appattempts") assert calls[2].args == ( "GET", f"{RM}/ws/v1/cluster/apps/application_1/appattempts/attempt_1/containers", ) assert calls[3].args == ( "GET", "http://nm1:8042/node/containerlogs/container_1/hdfs/", ) def test_logs_raises_on_404_when_nm_also_empty(): responses = [ _resp(404), # aggregated-logs H3 endpoint missing _resp(200, json_data={"appAttempts": {"appAttempt": []}}), ] with patch( "spark_executor.core.yarn_client.httpx.request", side_effect=responses ): with pytest.raises(YarnError, match="log-aggregation-enable"): get_application_logs("application_1", RM) def test_logs_handles_multiple_containers(): responses = [ _resp(404), # aggregated-logs H3 endpoint missing _resp(200, json_data={"appAttempts": {"appAttempt": [{"id": "attempt_1"}]}}), _resp( 200, json_data={ "containers": { "container": [ { "id": "container_1", "user": "hdfs", "nodeHttpAddress": "nm1:8042", }, { "id": "container_2", "user": "hdfs", "nodeHttpAddress": "nm2:8042", }, ] } }, ), _resp(200, text="log one"), _resp(200, text="log two"), ] with patch( "spark_executor.core.yarn_client.httpx.request", side_effect=responses ) as m: out = get_application_logs("application_1", RM) assert "log one" in out assert "log two" in out assert out.index("log one") < out.index("log two") calls = m.call_args_list assert calls[3].args == ( "GET", "http://nm1:8042/node/containerlogs/container_1/hdfs/", ) assert calls[4].args == ( "GET", "http://nm2:8042/node/containerlogs/container_2/hdfs/", ) def test_logs_5xx_on_aggregated_endpoint_raises_immediately(): with patch( "spark_executor.core.yarn_client.httpx.request", return_value=_resp(500, text="boom"), ) as m: with pytest.raises(YarnError, match="500"): get_application_logs("application_1", RM) assert m.call_count == 1 # --- kill_application --- def test_kill_sends_put_with_killed_state(): fake = _resp(200, json_data={"app": {"state": "KILLED"}}) with patch("spark_executor.core.yarn_client.httpx.request", return_value=fake) as m: kill_application("application_1", RM) args = m.call_args.args assert args == ("PUT", f"{RM}/ws/v1/cluster/apps/application_1/state") assert m.call_args.kwargs["json"] == {"state": "KILLED"} def test_kill_raises_on_5xx(): with patch( "spark_executor.core.yarn_client.httpx.request", return_value=_resp(403, text="forbidden"), ): with pytest.raises(YarnError, match="403"): kill_application("application_1", RM) # --- connection errors --- def test_status_wraps_httpx_errors_as_yarn_error(): with patch( "spark_executor.core.yarn_client.httpx.request", side_effect=httpx.ConnectError("connection refused"), ): with pytest.raises(YarnError, match="connection failed"): get_application_status("application_1", RM)