ClaudeandClaude
7d4e512cb0
- fetch_url: revalidate allowlist on every redirect hop (fixes SSRF where 302 to disallowed host / 169.254.169.254 / file:// bypassed the url_allowlist). Stream response body with iter_bytes and cap at 1MB so a multi-GB response from an allowlisted host cannot OOM the service. Reuses the manual-redirect-loop pattern from yarn_client. - list_applications: stop swallowing 404 (YARN returns 200+empty for "no match"; 404 means the RM doesn't support the endpoint — surface the YarnError instead of hiding it as an empty result). Add Field(ge=1, le=10000) to ListApplicationsRequest.limit so a runaway limit is rejected at the Pydantic layer with 422. - save_connection: PATCH semantics for existing records. Re-route to update_connection when the name already exists so partial updates (e.g. only master) no longer wipe url_allowlist back to []. Uses an _UNSET sentinel in the tool function to distinguish "omitted" from "None" without breaking the existing parameter list. - README: drop leading space on 5 new connection-tool table rows that was breaking GitHub Flavored Markdown table continuity. - Indentation: normalize connections.py and requests.py to 4-space indent (auth_password/auth_principal/auth_keytab were 3-space). Co-Authored-By: Claude <noreply@anthropic.com>
306 lines
10 KiB
Python
306 lines
10 KiB
Python
# coding=utf-8
|
|
import json
|
|
from unittest.mock import patch
|
|
|
|
import pytest
|
|
|
|
from spark_executor.core import connection_store
|
|
from spark_executor.core.yarn_client import YarnError
|
|
from spark_executor.models import ApplicationSummary, Connection
|
|
from spark_executor.tools import connections, external_jobs
|
|
from spark_executor.tools.requests import ListApplicationsRequest
|
|
|
|
|
|
def _fresh_stores():
|
|
"""Reset connection store singletons for a single test."""
|
|
store = connection_store.ConnectionStore()
|
|
connection_store.store = store
|
|
connections.store = store
|
|
external_jobs.conn_store = store
|
|
|
|
|
|
@pytest.fixture
|
|
def fresh_stores(tmp_path, monkeypatch):
|
|
"""Reset connection store singletons to an isolated tmp_path."""
|
|
monkeypatch.setattr(connection_store, "DEFAULT_DATA_DIR", str(tmp_path))
|
|
_fresh_stores()
|
|
|
|
|
|
def test_get_external_job_logs_returns_tailed():
|
|
_fresh_stores()
|
|
external_jobs.conn_store.save(
|
|
Connection(name="prod", master="yarn", yarn_rm_url="http://rm:8088")
|
|
)
|
|
long_log = "LOG" * 3000
|
|
with patch(
|
|
"spark_executor.tools.external_jobs.get_application_logs",
|
|
return_value=long_log,
|
|
) as m:
|
|
out = external_jobs.get_external_job_logs(
|
|
application_id="application_1", connection_name="prod", tail_chars=100
|
|
)
|
|
assert out == long_log[-100:]
|
|
args = m.call_args.args
|
|
assert args[0] == "application_1"
|
|
assert args[1].yarn_rm_url == "http://rm:8088"
|
|
|
|
|
|
def test_get_external_job_logs_returns_full_when_short():
|
|
_fresh_stores()
|
|
external_jobs.conn_store.save(
|
|
Connection(name="prod", master="yarn", yarn_rm_url="http://rm:8088")
|
|
)
|
|
short_log = "short log"
|
|
with patch(
|
|
"spark_executor.tools.external_jobs.get_application_logs",
|
|
return_value=short_log,
|
|
):
|
|
out = external_jobs.get_external_job_logs(
|
|
application_id="application_1", connection_name="prod", tail_chars=5000
|
|
)
|
|
assert out == short_log
|
|
|
|
|
|
def test_get_external_job_logs_raises_when_connection_missing():
|
|
_fresh_stores()
|
|
with pytest.raises(KeyError, match="Connection not found"):
|
|
external_jobs.get_external_job_logs(
|
|
application_id="application_1", connection_name="missing"
|
|
)
|
|
|
|
|
|
def test_get_external_job_status_returns_state():
|
|
_fresh_stores()
|
|
external_jobs.conn_store.save(
|
|
Connection(name="prod", master="yarn", yarn_rm_url="http://rm:8088")
|
|
)
|
|
raw = json.dumps({"app": {"state": "RUNNING"}})
|
|
with patch(
|
|
"spark_executor.tools.external_jobs.get_application_status",
|
|
return_value=("RUNNING", raw),
|
|
) as m:
|
|
out = external_jobs.get_external_job_status(
|
|
application_id="application_1", connection_name="prod"
|
|
)
|
|
assert out.application_id == "application_1"
|
|
assert out.state == "RUNNING"
|
|
assert out.raw == raw
|
|
args = m.call_args.args
|
|
assert args[0] == "application_1"
|
|
assert args[1].yarn_rm_url == "http://rm:8088"
|
|
|
|
|
|
def test_get_external_job_status_raises_when_connection_missing():
|
|
_fresh_stores()
|
|
with pytest.raises(KeyError, match="Connection not found"):
|
|
external_jobs.get_external_job_status(
|
|
application_id="application_1", connection_name="missing"
|
|
)
|
|
|
|
|
|
def test_get_external_job_result_parses_app_fields():
|
|
_fresh_stores()
|
|
external_jobs.conn_store.save(
|
|
Connection(name="prod", master="yarn", yarn_rm_url="http://rm:8088")
|
|
)
|
|
raw = json.dumps(
|
|
{
|
|
"app": {
|
|
"state": "FINISHED",
|
|
"finalStatus": "SUCCEEDED",
|
|
"diagnostics": "",
|
|
"trackingUrl": "http://rm:8088/proxy/application_1",
|
|
"startedTime": 100,
|
|
"finishedTime": 200,
|
|
}
|
|
}
|
|
)
|
|
with patch(
|
|
"spark_executor.tools.external_jobs.get_application_status",
|
|
return_value=("FINISHED", raw),
|
|
) as m:
|
|
out = external_jobs.get_external_job_result(
|
|
application_id="application_1", connection_name="prod"
|
|
)
|
|
assert out.application_id == "application_1"
|
|
assert out.state == "FINISHED"
|
|
assert out.final_status == "SUCCEEDED"
|
|
assert out.diagnostics == ""
|
|
assert out.tracking_url == "http://rm:8088/proxy/application_1"
|
|
assert out.started_time == 100
|
|
assert out.finished_time == 200
|
|
args = m.call_args.args
|
|
assert args[0] == "application_1"
|
|
assert args[1].yarn_rm_url == "http://rm:8088"
|
|
|
|
|
|
def test_get_external_job_result_raises_when_connection_missing():
|
|
_fresh_stores()
|
|
with pytest.raises(KeyError, match="Connection not found"):
|
|
external_jobs.get_external_job_result(
|
|
application_id="application_1", connection_name="missing"
|
|
)
|
|
|
|
|
|
def test_list_applications_returns_summaries(fresh_stores):
|
|
external_jobs.conn_store.save(
|
|
Connection(name="prod", master="yarn", yarn_rm_url="http://rm:8088")
|
|
)
|
|
with patch(
|
|
"spark_executor.tools.external_jobs.list_applications_yarn",
|
|
return_value=[
|
|
{
|
|
"id": "application_1",
|
|
"name": "app-one",
|
|
"user": "alice",
|
|
"queue": "default",
|
|
"state": "RUNNING",
|
|
"finalStatus": "UNDEFINED",
|
|
"applicationType": "SPARK",
|
|
"applicationTags": "tag1",
|
|
"startedTime": 1000,
|
|
"finishedTime": 0,
|
|
"trackingUrl": "http://rm:8088/proxy/application_1",
|
|
"progress": 75.0,
|
|
},
|
|
{
|
|
"id": "application_2",
|
|
"name": "app-two",
|
|
"user": "bob",
|
|
"queue": "research",
|
|
"state": "FINISHED",
|
|
"finalStatus": "SUCCEEDED",
|
|
"applicationType": "SPARK",
|
|
"applicationTags": "",
|
|
"startedTime": 2000,
|
|
"finishedTime": 3000,
|
|
"trackingUrl": "http://rm:8088/proxy/application_2",
|
|
"progress": 100.0,
|
|
},
|
|
],
|
|
) as m:
|
|
out = external_jobs.list_applications("prod")
|
|
assert len(out) == 2
|
|
assert all(isinstance(item, ApplicationSummary) for item in out)
|
|
assert out[0].application_id == "application_1"
|
|
assert out[1].application_id == "application_2"
|
|
args = m.call_args.args
|
|
assert args[0].yarn_rm_url == "http://rm:8088"
|
|
|
|
|
|
def test_list_applications_passes_state_filter(fresh_stores):
|
|
external_jobs.conn_store.save(
|
|
Connection(name="prod", master="yarn", yarn_rm_url="http://rm:8088")
|
|
)
|
|
with patch(
|
|
"spark_executor.tools.external_jobs.list_applications_yarn",
|
|
return_value=[],
|
|
) as m:
|
|
external_jobs.list_applications("prod", state="RUNNING")
|
|
assert m.call_args.kwargs == {"state": "RUNNING", "queue": None, "limit": 100}
|
|
|
|
|
|
def test_list_applications_passes_queue_filter(fresh_stores):
|
|
external_jobs.conn_store.save(
|
|
Connection(name="prod", master="yarn", yarn_rm_url="http://rm:8088")
|
|
)
|
|
with patch(
|
|
"spark_executor.tools.external_jobs.list_applications_yarn",
|
|
return_value=[],
|
|
) as m:
|
|
external_jobs.list_applications("prod", queue="research")
|
|
assert m.call_args.kwargs == {"state": None, "queue": "research", "limit": 100}
|
|
|
|
|
|
def test_list_applications_passes_limit_filter(fresh_stores):
|
|
external_jobs.conn_store.save(
|
|
Connection(name="prod", master="yarn", yarn_rm_url="http://rm:8088")
|
|
)
|
|
with patch(
|
|
"spark_executor.tools.external_jobs.list_applications_yarn",
|
|
return_value=[],
|
|
) as m:
|
|
external_jobs.list_applications("prod", limit=50)
|
|
assert m.call_args.kwargs == {"state": None, "queue": None, "limit": 50}
|
|
|
|
|
|
def test_list_applications_default_limit_is_100(fresh_stores):
|
|
external_jobs.conn_store.save(
|
|
Connection(name="prod", master="yarn", yarn_rm_url="http://rm:8088")
|
|
)
|
|
with patch(
|
|
"spark_executor.tools.external_jobs.list_applications_yarn",
|
|
return_value=[],
|
|
) as m:
|
|
external_jobs.list_applications("prod")
|
|
assert m.call_args.kwargs == {"state": None, "queue": None, "limit": 100}
|
|
assert ListApplicationsRequest(connection_name="prod").limit == 100
|
|
|
|
|
|
def test_list_applications_returns_empty_list_when_no_apps(fresh_stores):
|
|
external_jobs.conn_store.save(
|
|
Connection(name="prod", master="yarn", yarn_rm_url="http://rm:8088")
|
|
)
|
|
with patch(
|
|
"spark_executor.tools.external_jobs.list_applications_yarn",
|
|
return_value=[],
|
|
):
|
|
out = external_jobs.list_applications("prod")
|
|
assert out == []
|
|
|
|
|
|
def test_list_applications_raises_for_missing_connection(fresh_stores):
|
|
with pytest.raises(KeyError, match="Connection not found"):
|
|
external_jobs.list_applications("missing")
|
|
|
|
|
|
def test_list_applications_maps_yarn_json_to_summary(fresh_stores):
|
|
external_jobs.conn_store.save(
|
|
Connection(name="prod", master="yarn", yarn_rm_url="http://rm:8088")
|
|
)
|
|
yarn_app = {
|
|
"id": "application_42",
|
|
"name": "mapped-app",
|
|
"user": "carol",
|
|
"queue": "prod",
|
|
"state": "ACCEPTED",
|
|
"finalStatus": "UNDEFINED",
|
|
"applicationType": "MAPREDUCE",
|
|
"applicationTags": "batch",
|
|
"startedTime": 12345,
|
|
"finishedTime": 0,
|
|
"trackingUrl": "http://rm:8088/proxy/application_42",
|
|
"progress": 12.5,
|
|
}
|
|
with patch(
|
|
"spark_executor.tools.external_jobs.list_applications_yarn",
|
|
return_value=[yarn_app],
|
|
):
|
|
out = external_jobs.list_applications("prod")
|
|
assert len(out) == 1
|
|
summary = out[0]
|
|
assert summary.application_id == "application_42"
|
|
assert summary.name == "mapped-app"
|
|
assert summary.user == "carol"
|
|
assert summary.queue == "prod"
|
|
assert summary.state == "ACCEPTED"
|
|
assert summary.final_status == "UNDEFINED"
|
|
assert summary.application_type == "MAPREDUCE"
|
|
assert summary.application_tags == "batch"
|
|
assert summary.started_time == 12345
|
|
assert summary.finished_time == 0
|
|
assert summary.tracking_url == "http://rm:8088/proxy/application_42"
|
|
assert summary.progress == 12.5
|
|
|
|
|
|
def test_list_applications_raises_on_404(fresh_stores):
|
|
external_jobs.conn_store.save(
|
|
Connection(name="prod", master="yarn", yarn_rm_url="http://rm:8088")
|
|
)
|
|
with patch(
|
|
"spark_executor.tools.external_jobs.list_applications_yarn",
|
|
side_effect=YarnError("YARN list applications failed: 404"),
|
|
):
|
|
with pytest.raises(YarnError, match="YARN list applications failed"):
|
|
external_jobs.list_applications("prod")
|