diff --git a/README.md b/README.md index c644a6d..867c96a 100644 --- a/README.md +++ b/README.md @@ -100,11 +100,11 @@ MCP 客户端需要先执行 `initialize` 握手，拿到 `mcp-session-id` 后 | 工具 | 作用 | | --- | --- | - | `save_connection` | 保存或更新一个 Spark/YARN 连接配置 | - | `list_connections` | 列出所有连接配置 | - | `get_connection` | 按名称读取连接配置 | - | `delete_connection` | 删除连接配置 | - | `update_connection` | 部分更新一个已存在的连接 (PATCH 语义, 只改提供的字段) | +| `save_connection` | 保存或更新一个 Spark/YARN 连接配置 | +| `list_connections` | 列出所有连接配置 | +| `get_connection` | 按名称读取连接配置 | +| `delete_connection` | 删除连接配置 | +| `update_connection` | 部分更新一个已存在的连接 (PATCH 语义, 只改提供的字段) | | `write_job_file` | 将 LLM 已生成的 PySpark 代码写入服务端文件 | | `read_job_file` | 读取已存在的 PySpark 脚本内容 | | `update_job_file` | 覆盖更新已存在的 PySpark 脚本 | diff --git a/spark_executor/core/yarn_client.py b/spark_executor/core/yarn_client.py index cd45c5a..f86e3f5 100644 --- a/spark_executor/core/yarn_client.py +++ b/spark_executor/core/yarn_client.py @@ -299,9 +299,6 @@ def list_applications( resp = _request("GET", url, params=params, verify=config.verify_for_httpx(), auth=config.auth_for_httpx()) - if resp.status_code == 404: - # No apps match (or RM doesn't support the endpoint) - return [] if resp.status_code >= 400: raise YarnError( f"YARN list applications failed: {resp.status_code} {resp.text[:200]}" diff --git a/spark_executor/server.py b/spark_executor/server.py index 6453ca6..e43efb0 100644 --- a/spark_executor/server.py +++ b/spark_executor/server.py @@ -390,12 +390,20 @@ def _list_applications(req: ListApplicationsRequest): "yarn_rm_url is required for get_external_* to work** — spark-submit " "can discover the RM for submissions, but direct YARN REST queries " "need an explicit URL. Saving with an existing name overwrites the " - "record in place (no version history)." + "record in place (no version history). When the name already exists, " + "only the provided fields are changed (PATCH semantics); omitted " + "fields keep their previous values." ), ) def _save_connection(req: SaveConnectionRequest): - # exclude_none so we don't overwrite the function's default with explicit None - return save_connection(**req.model_dump(exclude_none=True)) + fields = req.model_dump(exclude_none=True) + name = fields.pop("name") + try: + get_connection(name) + except KeyError: + # exclude_none so we don't overwrite the function's default with explicit None + return save_connection(name=name, **fields) + return update_connection(name=name, **fields) @app.post( diff --git a/spark_executor/tools/connections.py b/spark_executor/tools/connections.py index 0af462b..4c631a7 100644 --- a/spark_executor/tools/connections.py +++ b/spark_executor/tools/connections.py @@ -4,45 +4,75 @@ @Author :tao.chen """ from common.logging import logger -from spark_executor.core.connection_store import ConnectionStore, store +from spark_executor.core.connection_store import store from spark_executor.models import Connection +_UNSET = object() + + def save_connection( *, name: str, master: str, - deploy_mode: str = "cluster", - yarn_rm_url: str | None = None, - spark_conf: dict[str, str] | None = None, - ssl_verify: bool | None = None, - ssl_ca_bundle: str | None = None, - auth_type: str = "none", - auth_user: str | None = None, - auth_password: str | None = None, - auth_principal: str | None = None, - auth_keytab: str | None = None, - url_allowlist: list[str] | None = None, -) -> dict[str, str]: + deploy_mode: str = _UNSET, # type: ignore[assignment] + yarn_rm_url: str | None = _UNSET, # type: ignore[assignment] + spark_conf: dict[str, str] | None = _UNSET, # type: ignore[assignment] + ssl_verify: bool | None = _UNSET, # type: ignore[assignment] + ssl_ca_bundle: str | None = _UNSET, # type: ignore[assignment] + auth_type: str = _UNSET, # type: ignore[assignment] + auth_user: str | None = _UNSET, # type: ignore[assignment] + auth_password: str | None = _UNSET, # type: ignore[assignment] + auth_principal: str | None = _UNSET, # type: ignore[assignment] + auth_keytab: str | None = _UNSET, # type: ignore[assignment] + url_allowlist: list[str] | None = _UNSET, # type: ignore[assignment] +) -> dict[str, object]: logger.debug( - f"save_connection enter name={name} master={master} deploy_mode={deploy_mode} " - f"yarn_rm_url={yarn_rm_url} spark_conf_keys={list((spark_conf or {}).keys())}" + f"save_connection enter name={name} master={master} " + f"spark_conf_keys={list((spark_conf if isinstance(spark_conf, dict) else {}).keys())}" ) - conn = Connection( - name=name, - master=master, - deploy_mode=deploy_mode, - yarn_rm_url=yarn_rm_url, - spark_conf=spark_conf or {}, - ssl_verify=ssl_verify, - ssl_ca_bundle=ssl_ca_bundle, - auth_type=auth_type, - auth_user=auth_user, - auth_password=auth_password, - auth_principal=auth_principal, - auth_keytab=auth_keytab, - url_allowlist=url_allowlist or [], - ) + existing = store.get(name) + if existing is not None: + fields: dict[str, object] = {"master": master} + if deploy_mode is not _UNSET: + fields["deploy_mode"] = deploy_mode + if yarn_rm_url is not _UNSET: + fields["yarn_rm_url"] = yarn_rm_url + if spark_conf is not _UNSET: + fields["spark_conf"] = spark_conf or {} + if ssl_verify is not _UNSET: + fields["ssl_verify"] = ssl_verify + if ssl_ca_bundle is not _UNSET: + fields["ssl_ca_bundle"] = ssl_ca_bundle + if auth_type is not _UNSET: + fields["auth_type"] = auth_type + if auth_user is not _UNSET: + fields["auth_user"] = auth_user + if auth_password is not _UNSET: + fields["auth_password"] = auth_password + if auth_principal is not _UNSET: + fields["auth_principal"] = auth_principal + if auth_keytab is not _UNSET: + fields["auth_keytab"] = auth_keytab + if url_allowlist is not _UNSET: + fields["url_allowlist"] = url_allowlist or [] + return update_connection(name, **fields) + new_fields: dict[str, object] = { + "name": name, + "master": master, + "deploy_mode": deploy_mode if deploy_mode is not _UNSET else "cluster", + "yarn_rm_url": yarn_rm_url if yarn_rm_url is not _UNSET else None, + "spark_conf": (spark_conf if spark_conf is not _UNSET else None) or {}, + "ssl_verify": ssl_verify if ssl_verify is not _UNSET else None, + "ssl_ca_bundle": ssl_ca_bundle if ssl_ca_bundle is not _UNSET else None, + "auth_type": auth_type if auth_type is not _UNSET else "none", + "auth_user": auth_user if auth_user is not _UNSET else None, + "auth_password": auth_password if auth_password is not _UNSET else None, + "auth_principal": auth_principal if auth_principal is not _UNSET else None, + "auth_keytab": auth_keytab if auth_keytab is not _UNSET else None, + "url_allowlist": (url_allowlist if url_allowlist is not _UNSET else None) or [], + } + conn = Connection(**new_fields) store.save(conn) return {"name": name, "status": "SAVED"} @@ -54,8 +84,8 @@ def update_connection(name: str, **fields) -> dict[str, object]: you pass are changed. To clear an optional field (e.g. `yarn_rm_url`), use `delete_connection(name=...)` followed by `save_connection(...)`. - Mutable fields: master, deploy_mode, yarn_rm_url, spark_conf, - ssl_verify, ssl_ca_bundle, auth_type, auth_user, auth_password, + Mutable fields: master, deploy_mode, yarn_rm_url, spark_conf, + ssl_verify, ssl_ca_bundle, auth_type, auth_user, auth_password, auth_principal, auth_keytab, url_allowlist. """ logger.debug(f"update_connection enter name={name} fields={sorted(fields.keys())}") diff --git a/spark_executor/tools/fetch_url.py b/spark_executor/tools/fetch_url.py index c505b88..4813c30 100644 --- a/spark_executor/tools/fetch_url.py +++ b/spark_executor/tools/fetch_url.py @@ -12,7 +12,7 @@ are intentionally NOT performed. Reuses the connection's saved auth/SSL config so the agent doesn't need cluster credentials. """ import fnmatch -from urllib.parse import urlparse +from urllib.parse import urlparse, urljoin import httpx @@ -24,6 +24,9 @@ from spark_executor.tools.connections import store as conn_store _MAX_BODY_BYTES = 1_000_000 # 1 MB cap on response body _REQUEST_TIMEOUT_SECONDS = 30 +_MAX_REDIRECTS = 3 +_REDIRECT_STATUSES = frozenset({301, 302, 303, 307, 308}) + def _host_matches_any_glob(host: str, allowlist: list[str]) -> bool: """True if `host` matches any of the fnmatch glob patterns. @@ -67,6 +70,29 @@ def _validate_url_host(url: str, allowlist: list[str] | None) -> None: ) +def _read_streamed_body(resp: httpx.Response) -> tuple[bytes, bool]: + """Stream the response body, keeping at most ``_MAX_BODY_BYTES`` bytes. + + Returns ``(body_bytes, truncated)``. Stops reading as soon as the cap is + exceeded so that a multi-gigabyte response from an allowlisted host cannot + OOM the service. + """ + chunks: list[bytes] = [] + total = 0 + truncated = False + for chunk in resp.iter_bytes(): + if total + len(chunk) <= _MAX_BODY_BYTES: + chunks.append(chunk) + total += len(chunk) + else: + if total < _MAX_BODY_BYTES: + chunks.append(chunk[: _MAX_BODY_BYTES - total]) + total = _MAX_BODY_BYTES + truncated = True + break + return b"".join(chunks), truncated + + def fetch_url(url: str, connection_name: str) -> FetchUrlResult: """Proxy an HTTP GET to url using the auth/SSL settings of connection_name.""" logger.debug(f"fetch_url enter url={url} connection_name={connection_name}") @@ -77,28 +103,50 @@ def fetch_url(url: str, connection_name: str) -> FetchUrlResult: _validate_url_host(url, conn.url_allowlist) config = YarnClientConfig.from_connection(conn) - resp = httpx.get( - url, - auth=config.auth_for_httpx(), - verify=config.verify_for_httpx(), - timeout=_REQUEST_TIMEOUT_SECONDS, - follow_redirects=True, - ) + auth = config.auth_for_httpx() + verify = config.verify_for_httpx() - body = resp.text[:_MAX_BODY_BYTES] - truncated = len(resp.text) > _MAX_BODY_BYTES + for hop in range(_MAX_REDIRECTS + 1): + with httpx.stream( + "GET", + url, + auth=auth, + verify=verify, + timeout=_REQUEST_TIMEOUT_SECONDS, + follow_redirects=False, + ) as resp: + if resp.status_code not in _REDIRECT_STATUSES: + body_bytes, truncated = _read_streamed_body(resp) + encoding = resp.encoding or "utf-8" + body = body_bytes.decode(encoding, errors="replace") - logger.info( - f"fetch_url ok url={url} connection_name={connection_name} " - f"status_code={resp.status_code} " - f"content_type={resp.headers.get('content-type', '')} " - f"body_bytes={len(body)} truncated={truncated}" - ) + logger.info( + f"fetch_url ok url={url} connection_name={connection_name} " + f"status_code={resp.status_code} " + f"content_type={resp.headers.get('content-type', '')} " + f"body_bytes={len(body)} truncated={truncated}" + ) - return FetchUrlResult( - url=url, - status_code=resp.status_code, - content_type=resp.headers.get("content-type", ""), - body=body, - truncated=truncated, + return FetchUrlResult( + url=url, + status_code=resp.status_code, + content_type=resp.headers.get("content-type", ""), + body=body, + truncated=truncated, + ) + + location = resp.headers.get("Location") or resp.headers.get("location") + if not location: + raise ValueError( + f"HTTP GET returned {resp.status_code} with no Location header at {url}" + ) + next_url = urljoin(url, location) + logger.debug( + f"fetch_url redirect url={url} status={resp.status_code} to={next_url}" + ) + _validate_url_host(next_url, conn.url_allowlist) + url = next_url + + raise ValueError( + f"fetch_url exceeded {_MAX_REDIRECTS} redirects (last url: {url})" ) diff --git a/spark_executor/tools/requests.py b/spark_executor/tools/requests.py index 27f8284..9cb1f0f 100644 --- a/spark_executor/tools/requests.py +++ b/spark_executor/tools/requests.py @@ -117,10 +117,10 @@ class SaveConnectionRequest(BaseModel): description=( "Absolute path to a Kerberos keytab file. Optional convenience " "for 'kinit -kt' workflows. The service does NOT auto-initialize " - "from the keytab — you must `kinit -kt ` " - "yourself before calling the tools." - ), - ) + "from the keytab — you must `kinit -kt ` " + "yourself before calling the tools." + ), + ) url_allowlist: list[str] | None = Field( default=None, @@ -504,6 +504,8 @@ class ListApplicationsRequest(BaseModel): ) limit: int = Field( default=100, + ge=1, + le=10000, description=( "Maximum number of applications to return. YARN has no " "offset-based pagination, so for large clusters use state/queue " diff --git a/tests/unit/test_connection_tools.py b/tests/unit/test_connection_tools.py index 933c46b..dad74c7 100644 --- a/tests/unit/test_connection_tools.py +++ b/tests/unit/test_connection_tools.py @@ -121,6 +121,15 @@ def test_delete_connection_returns_status(_fresh_store): def test_delete_connection_unknown_raises(_fresh_store): with pytest.raises(KeyError): connections.delete_connection("missing") + + +def test_save_connection_preserves_url_allowlist_on_existing_record(_fresh_store): + connections.save_connection(name="prod", master="yarn", url_allowlist=["ccam*"]) + out = connections.save_connection(name="prod", master="spark://new:7077") + assert out["master"] == "spark://new:7077" + assert out["url_allowlist"] == ["ccam*"] + + def test_update_connection_changes_specified_field(_fresh_store): connections.save_connection(name="prod", master="yarn") out = connections.update_connection(name="prod", master="spark://new:7077") diff --git a/tests/unit/test_external_jobs.py b/tests/unit/test_external_jobs.py index 36e3f11..7b22e21 100644 --- a/tests/unit/test_external_jobs.py +++ b/tests/unit/test_external_jobs.py @@ -5,6 +5,7 @@ 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 @@ -290,3 +291,15 @@ def test_list_applications_maps_yarn_json_to_summary(fresh_stores): 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") diff --git a/tests/unit/test_fetch_url.py b/tests/unit/test_fetch_url.py index a74d3ac..33d556d 100644 --- a/tests/unit/test_fetch_url.py +++ b/tests/unit/test_fetch_url.py @@ -4,11 +4,25 @@ from unittest.mock import patch import httpx import pytest +from pydantic import ValidationError from spark_executor.core import connection_store from spark_executor.core.connection_store import ConnectionStore from spark_executor.models import Connection from spark_executor.tools import connections, fetch_url +from spark_executor.tools.requests import ListApplicationsRequest + + +def _stream_cm(resp): + """Wrap a response (or fake response) in a context manager for httpx.stream.""" + class _CM: + def __enter__(self): + return resp + + def __exit__(self, exc_type, exc, tb): + return False + + return _CM() def _fresh_stores(tmp_path, monkeypatch): @@ -37,7 +51,9 @@ def test_fetch_url_returns_body_and_status(fresh_stores): ) ) resp = httpx.Response(200, text="hello", headers={"content-type": "text/html"}) - with patch("spark_executor.tools.fetch_url.httpx.get", return_value=resp) as m: + with patch( + "spark_executor.tools.fetch_url.httpx.stream", return_value=_stream_cm(resp) + ) as m: out = fetch_url.fetch_url("http://nm01.prod.internal:8042/node", "prod") assert out.url == "http://nm01.prod.internal:8042/node" assert out.status_code == 200 @@ -63,7 +79,9 @@ def test_fetch_url_truncates_body_over_1mb(fresh_stores): ) big = "x" * (fetch_url._MAX_BODY_BYTES + 1) resp = httpx.Response(200, text=big) - with patch("spark_executor.tools.fetch_url.httpx.get", return_value=resp): + with patch( + "spark_executor.tools.fetch_url.httpx.stream", return_value=_stream_cm(resp) + ): out = fetch_url.fetch_url("http://nm01.prod.internal/big", "prod") assert out.truncated is True assert len(out.body) == fetch_url._MAX_BODY_BYTES @@ -82,7 +100,9 @@ def test_fetch_url_passes_auth_from_connection(fresh_stores): ) ) resp = httpx.Response(200, text="ok") - with patch("spark_executor.tools.fetch_url.httpx.get", return_value=resp) as m: + with patch( + "spark_executor.tools.fetch_url.httpx.stream", return_value=_stream_cm(resp) + ) as m: fetch_url.fetch_url("http://nm01.prod.internal:8042/node", "auth") auth = m.call_args.kwargs["auth"] assert isinstance(auth, httpx.BasicAuth) @@ -103,26 +123,13 @@ def test_fetch_url_passes_ssl_verify_from_connection(fresh_stores): ) ) resp = httpx.Response(200, text="ok") - with patch("spark_executor.tools.fetch_url.httpx.get", return_value=resp) as m: + with patch( + "spark_executor.tools.fetch_url.httpx.stream", return_value=_stream_cm(resp) + ) as m: fetch_url.fetch_url("http://nm01.prod.internal:8042/node", "insecure") assert m.call_args.kwargs["verify"] is False -def test_fetch_url_follows_redirects(fresh_stores): - fetch_url.conn_store.save( - Connection( - name="prod", - master="yarn", - yarn_rm_url="http://rm.prod.internal:8088", - url_allowlist=["*.prod.internal"], - ) - ) - resp = httpx.Response(200, text="ok") - with patch("spark_executor.tools.fetch_url.httpx.get", return_value=resp) as m: - fetch_url.fetch_url("http://nm01.prod.internal:8042/node", "prod") - assert m.call_args.kwargs["follow_redirects"] is True - - def test_fetch_url_allows_host_matching_glob_pattern(fresh_stores): fetch_url.conn_store.save( Connection( @@ -133,7 +140,9 @@ def test_fetch_url_allows_host_matching_glob_pattern(fresh_stores): ) ) resp = httpx.Response(200, text="hello") - with patch("spark_executor.tools.fetch_url.httpx.get", return_value=resp) as m: + with patch( + "spark_executor.tools.fetch_url.httpx.stream", return_value=_stream_cm(resp) + ) as m: out = fetch_url.fetch_url("http://ccam50:8088/foo", "ccam") assert out.status_code == 200 assert out.body == "hello" @@ -150,7 +159,9 @@ def test_fetch_url_allows_host_matching_any_of_multiple_globs(fresh_stores): ) ) resp = httpx.Response(200, text="hello") - with patch("spark_executor.tools.fetch_url.httpx.get", return_value=resp) as m: + with patch( + "spark_executor.tools.fetch_url.httpx.stream", return_value=_stream_cm(resp) + ) as m: out = fetch_url.fetch_url( "http://history.prod.internal:18080/api/v1/info", "prod" ) @@ -195,7 +206,9 @@ def test_fetch_url_glob_match_does_not_require_suffix_overlap(fresh_stores): ) ) resp = httpx.Response(200, text="hello") - with patch("spark_executor.tools.fetch_url.httpx.get", return_value=resp) as m: + with patch( + "spark_executor.tools.fetch_url.httpx.stream", return_value=_stream_cm(resp) + ) as m: out = fetch_url.fetch_url("http://ccam50:8088/", "prod") assert out.status_code == 200 assert out.body == "hello" @@ -211,7 +224,9 @@ def test_fetch_url_accepts_ip_literal_when_in_url_allowlist(fresh_stores): ) ) resp = httpx.Response(200, text="hello") - with patch("spark_executor.tools.fetch_url.httpx.get", return_value=resp) as m: + with patch( + "spark_executor.tools.fetch_url.httpx.stream", return_value=_stream_cm(resp) + ) as m: out = fetch_url.fetch_url("http://10.0.0.1/secret", "prod") assert out.status_code == 200 assert out.body == "hello" @@ -227,7 +242,9 @@ def test_fetch_url_accepts_https_when_in_url_allowlist(fresh_stores): ) ) resp = httpx.Response(200, text="hello") - with patch("spark_executor.tools.fetch_url.httpx.get", return_value=resp) as m: + with patch( + "spark_executor.tools.fetch_url.httpx.stream", return_value=_stream_cm(resp) + ) as m: out = fetch_url.fetch_url("https://ccam1.example.com/secure", "prod") assert out.status_code == 200 assert out.body == "hello" @@ -271,3 +288,117 @@ def test_fetch_url_omitted_url_allowlist_defaults_to_empty_and_rejects(fresh_sto ) with pytest.raises(ValueError, match="is not in Connection.url_allowlist"): fetch_url.fetch_url("http://nm.prod.internal/", "prod") + + +def test_fetch_url_rejects_redirect_to_disallowed_host(fresh_stores): + fetch_url.conn_store.save( + Connection( + name="ccam", + master="yarn", + yarn_rm_url="http://ccam1:8088", + url_allowlist=["ccam*"], + ) + ) + redirect = httpx.Response(302, headers={"Location": "http://evil.com/"}) + with patch( + "spark_executor.tools.fetch_url.httpx.stream", + return_value=_stream_cm(redirect), + ) as m: + with pytest.raises(ValueError, match="is not in Connection.url_allowlist"): + fetch_url.fetch_url("http://ccam50/foo", "ccam") + assert m.call_count == 1 + + +def test_fetch_url_follows_redirect_to_allowed_host(fresh_stores): + fetch_url.conn_store.save( + Connection( + name="prod", + master="yarn", + yarn_rm_url="http://rm.prod.internal:8088", + url_allowlist=["*.internal"], + ) + ) + redirect = httpx.Response( + 302, headers={"Location": "http://other.internal/"} + ) + final = httpx.Response(200, text="ok") + with patch( + "spark_executor.tools.fetch_url.httpx.stream", + side_effect=[_stream_cm(redirect), _stream_cm(final)], + ) as m: + out = fetch_url.fetch_url("http://foo.internal/", "prod") + assert out.status_code == 200 + assert out.body == "ok" + assert m.call_count == 2 + assert m.call_args_list[1].args[1] == "http://other.internal/" + + +def test_fetch_url_rejects_redirect_to_ip_literal_not_in_allowlist(fresh_stores): + fetch_url.conn_store.save( + Connection( + name="ccam", + master="yarn", + yarn_rm_url="http://ccam1:8088", + url_allowlist=["ccam*"], + ) + ) + redirect = httpx.Response(302, headers={"Location": "http://10.0.0.1/"}) + with patch( + "spark_executor.tools.fetch_url.httpx.stream", + return_value=_stream_cm(redirect), + ) as m: + with pytest.raises(ValueError, match="is not in Connection.url_allowlist"): + fetch_url.fetch_url("http://ccam50/foo", "ccam") + assert m.call_count == 1 + + +def test_fetch_url_streams_body_and_truncates_without_buffering_full(fresh_stores): + """The response body must be read incrementally; .text must not be accessed.""" + + class _FakeStreamResponse: + def __init__(self, chunks): + self.status_code = 200 + self.headers = httpx.Headers({"content-type": "text/plain"}) + self.encoding = "utf-8" + self._chunks = chunks + + def iter_bytes(self): + yield from self._chunks + + @property + def text(self): + raise AssertionError("response.text should not be accessed") + + fetch_url.conn_store.save( + Connection( + name="prod", + master="yarn", + url_allowlist=["*.prod.internal"], + ) + ) + chunk_size = 400_000 + chunks = [ + b"a" * chunk_size, + b"b" * chunk_size, + b"c" * chunk_size, + ] + fake = _FakeStreamResponse(chunks) + with patch( + "spark_executor.tools.fetch_url.httpx.stream", + return_value=_stream_cm(fake), + ) as m: + out = fetch_url.fetch_url("http://nm01.prod.internal/big", "prod") + assert m.call_count == 1 + assert out.truncated is True + assert len(out.body) == fetch_url._MAX_BODY_BYTES + assert out.body.startswith("a" * chunk_size) + # Only the first 200 KB of the third chunk were consumed before truncation. + assert out.body[-200_000:] == "c" * 200_000 + + +def test_list_applications_request_limit_bounds(): + assert ListApplicationsRequest(connection_name="prod", limit=10000).limit == 10000 + with pytest.raises(ValidationError): + ListApplicationsRequest(connection_name="prod", limit=0) + with pytest.raises(ValidationError): + ListApplicationsRequest(connection_name="prod", limit=10001)