# Copyright (c) Microsoft. All rights reserved. """ Tests for RESTful API pagination, sorting, and filtering functionality. Test categories: - Rollouts pagination, sorting, and filtering - Attempts pagination and sorting - Resources pagination and sorting - Spans pagination, sorting, and filtering """ import asyncio import contextlib from typing import Any, AsyncGenerator, Dict, List, Tuple import aiohttp import pytest import pytest_asyncio from portpicker import pick_unused_port from agentlightning.store import LightningStore from agentlightning.store.client_server import LightningStoreClient, LightningStoreServer from agentlightning.store.memory import InMemoryLightningStore from agentlightning.types import ( LLM, AttemptedRollout, OtelResource, PaginatedResult, PromptTemplate, Rollout, Span, TraceStatus, ) pytestmark = [pytest.mark.store] def _make_span(rollout_id: str, attempt_id: str, sequence_id: int, name: str) -> Span: return Span( rollout_id=rollout_id, attempt_id=attempt_id, sequence_id=sequence_id, trace_id=f"{sequence_id:032x}", span_id=f"{sequence_id:016x}", parent_id=None, name=name, status=TraceStatus(status_code="OK"), attributes={}, events=[], links=[], start_time=1.0, end_time=2.0, context=None, parent=None, resource=OtelResource(attributes={}, schema_url=""), ) @contextlib.asynccontextmanager async def _run_server_with_cors(cors_origins: List[str] | str | None = None): store = InMemoryLightningStore() port = pick_unused_port() server = LightningStoreServer(store, "127.0.0.1", port, cors_allow_origins=cors_origins) await server.start() session = aiohttp.ClientSession() try: yield server, session finally: await session.close() await server.stop() async def _request_rollouts_page( session: aiohttp.ClientSession, api_endpoint: str, method: str, payload: Dict[str, Any] ) -> Dict[str, Any]: """Send either GET or POST /search request for rollouts and return parsed payload.""" if method == "get": async with session.get(f"{api_endpoint}/rollouts", params=payload) as resp: assert resp.status == 200 return await resp.json() async with session.post(f"{api_endpoint}/rollouts/search", json=payload) as resp: assert resp.status == 200 return await resp.json() async def _request_attempts_page( session: aiohttp.ClientSession, api_endpoint: str, rollout_id: str, method: str, payload: Dict[str, Any], ) -> Dict[str, Any]: """Send either GET or POST /search request for attempts and return parsed payload.""" base = f"{api_endpoint}/rollouts/{rollout_id}/attempts" if method == "get": async with session.get(base, params=payload) as resp: assert resp.status == 200 return await resp.json() async with session.post(f"{base}/search", json=payload) as resp: assert resp.status == 200 return await resp.json() async def _request_spans_page( session: aiohttp.ClientSession, api_endpoint: str, rollout_id: str, method: str, payload: Dict[str, Any], ) -> Dict[str, Any]: """Send either GET or POST /search request for spans and return parsed payload.""" base_params = {"rollout_id": rollout_id} base_params.update(payload) if method == "get": async with session.get(f"{api_endpoint}/spans", params=base_params) as resp: assert resp.status == 200 return await resp.json() async with session.post(f"{api_endpoint}/spans/search", json=base_params) as resp: assert resp.status == 200 return await resp.json() async def _request_workers_page( session: aiohttp.ClientSession, api_endpoint: str, method: str, payload: Dict[str, Any], ) -> Dict[str, Any]: """Send either GET or POST /search request for workers and return parsed payload.""" base = f"{api_endpoint}/workers" if method == "get": async with session.get(base, params=payload) as resp: assert resp.status == 200 return await resp.json() async with session.post(f"{base}/search", json=payload) as resp: assert resp.status == 200 return await resp.json() @pytest_asyncio.fixture async def server_client( store_fixture: LightningStore, ) -> AsyncGenerator[Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str], None]: port = pick_unused_port() server = LightningStoreServer(store_fixture, "127.0.0.1", port) await server.start() client = LightningStoreClient(server.endpoint) session = aiohttp.ClientSession() # Get the full API endpoint with /v1/agl prefix api_endpoint = client.server_address try: yield server, client, session, api_endpoint finally: await session.close() await client.close() await server.stop() # CORS configuration tests @pytest.mark.asyncio async def test_cors_allows_specific_origin() -> None: async with _run_server_with_cors(cors_origins=["http://localhost:3000"]) as (server, session): url = f"{server.endpoint}/v1/agl/health" origin = "http://localhost:3000" async with session.get(url, headers={"Origin": origin}) as resp: assert resp.status == 200 assert resp.headers.get("access-control-allow-origin") == origin assert resp.headers.get("access-control-allow-credentials") == "true" async with session.options( url, headers={ "Origin": origin, "Access-Control-Request-Method": "GET", }, ) as resp: assert resp.status == 200 assert resp.headers.get("access-control-allow-origin") == origin allow_methods = resp.headers.get("access-control-allow-methods") or "" assert "GET" in {method.strip() for method in allow_methods.split(",") if method} @pytest.mark.asyncio async def test_cors_disallows_unconfigured_origin() -> None: async with _run_server_with_cors(cors_origins=["http://localhost:3000"]) as (server, session): url = f"{server.endpoint}/v1/agl/health" async with session.get(url, headers={"Origin": "http://malicious.example"}) as resp: assert resp.status == 200 assert "access-control-allow-origin" not in resp.headers @pytest.mark.asyncio async def test_cors_allows_wildcard_origin() -> None: async with _run_server_with_cors(cors_origins="*") as (server, session): url = f"{server.endpoint}/v1/agl/health" origin = "https://wildcard.example" async with session.get(url, headers={"Origin": origin}) as resp: assert resp.status == 200 allow_origin = resp.headers.get("access-control-allow-origin") assert allow_origin in {origin, "*"} allow_credentials = resp.headers.get("access-control-allow-credentials") if allow_credentials is not None: assert allow_credentials == "true" @pytest.mark.asyncio async def test_enqueue_endpoint_batches_payloads( server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str], ) -> None: _server, _client, session, api_endpoint = server_client single_payload = {"rollouts": [{"input": {"task": "single"}}]} async with session.post(f"{api_endpoint}/queues/rollouts/enqueue", json=single_payload) as resp: assert resp.status == 201 body = await resp.json() assert isinstance(body, list) assert len(body) == 1 # type: ignore assert body[0]["input"] == {"task": "single"} batch_payload = { "rollouts": [ {"input": {"task": "batch-1"}, "metadata": {"batch": 1}}, {"input": {"task": "batch-2"}}, ] } async with session.post(f"{api_endpoint}/queues/rollouts/enqueue", json=batch_payload) as resp: assert resp.status == 201 body = await resp.json() assert len(body) == 2 assert [item["input"]["task"] for item in body] == ["batch-1", "batch-2"] assert body[0]["metadata"] == {"batch": 1} @pytest.mark.asyncio async def test_dequeue_endpoint_returns_batches( server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str], ) -> None: server, _client, session, api_endpoint = server_client for idx in range(3): await server.enqueue_rollout(input={"idx": idx}) async with session.post( f"{api_endpoint}/queues/rollouts/dequeue", json={"limit": 2, "worker_id": "rest-worker"} ) as resp: assert resp.status == 200 body = await resp.json() assert len(body) == 2 assert all(item["attempt"]["worker_id"] == "rest-worker" for item in body) async with session.post(f"{api_endpoint}/queues/rollouts/dequeue") as resp: assert resp.status == 200 body = await resp.json() assert len(body) == 1 async with session.post(f"{api_endpoint}/queues/rollouts/dequeue") as resp: assert resp.status == 200 body = await resp.json() assert body == [] @pytest.mark.asyncio async def test_enqueue_endpoint_requires_rollouts_field( server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str], ) -> None: _server, _client, session, api_endpoint = server_client async with session.post(f"{api_endpoint}/queues/rollouts/enqueue", json={}) as resp: assert resp.status == 422 @pytest.mark.asyncio async def test_dequeue_endpoint_zero_limit_returns_empty( server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str], ) -> None: server, _client, session, api_endpoint = server_client await server.enqueue_rollout(input={"idx": 0}) async with session.post(f"{api_endpoint}/queues/rollouts/dequeue", json={"limit": 0}) as resp: assert resp.status == 200 assert await resp.json() == [] async with session.post(f"{api_endpoint}/queues/rollouts/dequeue", json={"limit": 1}) as resp: assert resp.status == 200 body = await resp.json() assert len(body) == 1 @pytest.mark.asyncio async def test_statistics_endpoint_returns_counts( server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str], ) -> None: """The statistics endpoint should expose store-level metrics.""" server, client, session, api_endpoint = server_client await client.start_rollout(input={"source": "stats-endpoint"}) async with session.get(f"{api_endpoint}/statistics") as resp: assert resp.status == 200 payload = await resp.json() expected_name = server.store.__class__.__name__ if server.store is not None else payload["name"] assert payload["name"] == expected_name assert payload["total_rollouts"] >= 1 @pytest.mark.asyncio async def test_rest_start_rollout_propagates_worker_id( server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str], ) -> None: server, _client, session, api_endpoint = server_client payload = {"input": {"source": "rest-worker"}, "worker_id": "rest-start-worker"} async with session.post(f"{api_endpoint}/rollouts", json=payload) as resp: assert resp.status == 201 data = await resp.json() assert data["attempt"]["worker_id"] == "rest-start-worker" worker = await server.get_worker_by_id("rest-start-worker") assert worker is not None assert worker.status == "busy" assert worker.current_rollout_id == data["rollout_id"] assert worker.current_attempt_id == data["attempt"]["attempt_id"] @pytest.mark.asyncio async def test_rest_start_attempt_propagates_worker_id( server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str], ) -> None: server, _client, session, api_endpoint = server_client async with session.post(f"{api_endpoint}/rollouts", json={"input": {"source": "rest-retry"}}) as resp: assert resp.status == 201 base_rollout = await resp.json() attempt_worker = "rest-attempt-worker" async with session.post( f"{api_endpoint}/rollouts/{base_rollout['rollout_id']}/attempts", json={"worker_id": attempt_worker}, ) as resp: assert resp.status == 201 retry_payload = await resp.json() assert retry_payload["attempt"]["sequence_id"] == 2 assert retry_payload["attempt"]["worker_id"] == attempt_worker worker = await server.get_worker_by_id(attempt_worker) assert worker is not None assert worker.status == "busy" assert worker.current_rollout_id == retry_payload["rollout_id"] assert worker.current_attempt_id == retry_payload["attempt"]["attempt_id"] # Rollouts Pagination, Sorting, and Filtering Tests @pytest.mark.asyncio async def test_rollouts_pagination_basic( server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str], ) -> None: """Test basic pagination for rollouts endpoint.""" server, _client, session, api_endpoint = server_client # Create 10 rollouts for i in range(10): await server.enqueue_rollout(input={"index": i}) # Get first page with limit=3 async with session.get(f"{api_endpoint}/rollouts", params={"limit": 3, "offset": 0}) as resp: assert resp.status == 200 data = await resp.json() assert data["total"] == 10 assert data["limit"] == 3 assert data["offset"] == 0 assert len(data["items"]) == 3 # Get second page async with session.get(f"{api_endpoint}/rollouts", params={"limit": 3, "offset": 3}) as resp: assert resp.status == 200 data = await resp.json() assert data["total"] == 10 assert data["limit"] == 3 assert data["offset"] == 3 assert len(data["items"]) == 3 @pytest.mark.asyncio async def test_rollouts_pagination_disabled( server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str], ) -> None: """Test pagination can be disabled with limit=-1.""" server, _client, session, api_endpoint = server_client # Create 15 rollouts for i in range(15): await server.enqueue_rollout(input={"index": i}) # Get all rollouts with limit=-1 async with session.get(f"{api_endpoint}/rollouts", params={"limit": -1}) as resp: assert resp.status == 200 data = await resp.json() assert data["total"] == 15 assert data["limit"] == -1 assert len(data["items"]) == 15 @pytest.mark.asyncio @pytest.mark.parametrize("method", ["get", "post"]) async def test_rollouts_search_supports_get_and_post( server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str], method: str ) -> None: """Ensure both GET and POST /search endpoints behave the same for rollouts.""" server, _client, session, api_endpoint = server_client for i in range(3): await server.enqueue_rollout(input={"index": i}) data = await _request_rollouts_page(session, api_endpoint, method, {"limit": 2, "offset": 0}) assert data["total"] == 3 assert data["limit"] == 2 assert data["offset"] == 0 assert len(data["items"]) == 2 @pytest.mark.asyncio async def test_rollouts_sorting_by_start_time( server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str], ) -> None: """Test sorting rollouts by start_time.""" server, _client, session, api_endpoint = server_client # Create rollouts with small delays rollouts: List[Rollout] = [] for i in range(5): r = await server.enqueue_rollout(input={"index": i}) rollouts.append(r) await asyncio.sleep(0.01) # Sort ascending by start_time async with session.get( f"{api_endpoint}/rollouts", params={"sort_by": "start_time", "sort_order": "asc", "limit": -1} ) as resp: assert resp.status == 200 data = await resp.json() items = data["items"] assert len(items) == 5 # Should be in ascending order for i in range(len(items) - 1): assert items[i]["start_time"] <= items[i + 1]["start_time"] # Sort descending by start_time (default) async with session.get( f"{api_endpoint}/rollouts", params={"sort_by": "start_time", "sort_order": "desc", "limit": -1} ) as resp: assert resp.status == 200 data = await resp.json() items = data["items"] # Should be in descending order for i in range(len(items) - 1): assert items[i]["start_time"] >= items[i + 1]["start_time"] @pytest.mark.asyncio async def test_rollouts_filter_by_status( server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str], ) -> None: """Test filtering rollouts by status.""" server, _client, session, api_endpoint = server_client # Create rollouts with different statuses r1 = await server.enqueue_rollout(input={"id": 1}) r2 = await server.enqueue_rollout(input={"id": 2}) _r3 = await server.enqueue_rollout(input={"id": 3}) await server.update_rollout(rollout_id=r1.rollout_id, status="succeeded") await server.update_rollout(rollout_id=r2.rollout_id, status="failed") # r3 remains queuing # Filter by single status async with session.get(f"{api_endpoint}/rollouts", params={"status_in": ["succeeded"], "limit": -1}) as resp: assert resp.status == 200 data = await resp.json() assert data["total"] == 1 assert data["items"][0]["rollout_id"] == r1.rollout_id # Filter by multiple statuses async with session.get( f"{api_endpoint}/rollouts", params={"status_in": ["succeeded", "failed"], "limit": -1} ) as resp: assert resp.status == 200 data = await resp.json() assert data["total"] == 2 rollout_ids = {item["rollout_id"] for item in data["items"]} assert rollout_ids == {r1.rollout_id, r2.rollout_id} @pytest.mark.asyncio async def test_rollouts_filter_by_rollout_id_in( server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str], ) -> None: """Test filtering rollouts by rollout_id_in.""" server, _client, session, api_endpoint = server_client # Create multiple rollouts r1 = await server.enqueue_rollout(input={"id": 1}) _r2 = await server.enqueue_rollout(input={"id": 2}) r3 = await server.enqueue_rollout(input={"id": 3}) # Filter by specific rollout IDs async with session.get( f"{api_endpoint}/rollouts", params={"rollout_id_in": [r1.rollout_id, r3.rollout_id], "limit": -1} ) as resp: assert resp.status == 200 data = await resp.json() assert data["total"] == 2 rollout_ids = {item["rollout_id"] for item in data["items"]} assert rollout_ids == {r1.rollout_id, r3.rollout_id} @pytest.mark.asyncio async def test_rollouts_filter_by_rollout_id_contains( server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str], ) -> None: """Test filtering rollouts by rollout_id_contains.""" server, _client, session, api_endpoint = server_client # Create rollouts r1 = await server.enqueue_rollout(input={"id": 1}) _r2 = await server.enqueue_rollout(input={"id": 2}) _r3 = await server.enqueue_rollout(input={"id": 3}) # Extract a substring from r1's ID substring = r1.rollout_id[3:8] # Get middle part of ID # Filter by substring async with session.get(f"{api_endpoint}/rollouts", params={"rollout_id_contains": substring, "limit": -1}) as resp: assert resp.status == 200 data = await resp.json() assert data["total"] >= 1 # Verify all results contain the substring for item in data["items"]: assert substring in item["rollout_id"] @pytest.mark.asyncio async def test_rollouts_filter_logic_and( server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str], ) -> None: """Test rollouts filtering with AND logic.""" server, _client, session, api_endpoint = server_client # Create rollouts r1 = await server.enqueue_rollout(input={"id": 1}) r2 = await server.enqueue_rollout(input={"id": 2}) r3 = await server.enqueue_rollout(input={"id": 3}) await server.update_rollout(rollout_id=r1.rollout_id, status="succeeded") await server.update_rollout(rollout_id=r2.rollout_id, status="succeeded") await server.update_rollout(rollout_id=r3.rollout_id, status="failed") # Filter with AND logic: status=succeeded AND rollout_id in list async with session.get( f"{api_endpoint}/rollouts", params={ "status_in": ["succeeded"], "rollout_id_in": [r1.rollout_id, r3.rollout_id], "filter_logic": "and", "limit": -1, }, ) as resp: assert resp.status == 200 data = await resp.json() # Only r1 matches both conditions (r3 is failed, r2 is not in the ID list) assert data["total"] == 1 assert data["items"][0]["rollout_id"] == r1.rollout_id @pytest.mark.asyncio async def test_rollouts_filter_logic_or( server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str], ) -> None: """Test rollouts filtering with OR logic.""" server, _client, session, api_endpoint = server_client # Create rollouts r1 = await server.enqueue_rollout(input={"id": 1}) _r2 = await server.enqueue_rollout(input={"id": 2}) r3 = await server.enqueue_rollout(input={"id": 3}) await server.update_rollout(rollout_id=r1.rollout_id, status="succeeded") # r2 and r3 remain queuing # Filter with OR logic: status=succeeded OR rollout_id=r3 async with session.get( f"{api_endpoint}/rollouts", params={"status_in": ["succeeded"], "rollout_id_in": [r3.rollout_id], "filter_logic": "or", "limit": -1}, ) as resp: assert resp.status == 200 data = await resp.json() # Both r1 (succeeded) and r3 (in ID list) should match assert data["total"] == 2 rollout_ids = {item["rollout_id"] for item in data["items"]} assert rollout_ids == {r1.rollout_id, r3.rollout_id} @pytest.mark.asyncio async def test_rollouts_sorting_with_none_values( server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str], ) -> None: """Test sorting rollouts by fields that may have None values.""" server, _client, session, api_endpoint = server_client # Create rollouts - mode is optional and can be None _r1 = await server.enqueue_rollout(input={"id": 1}, mode="train") r2 = await server.enqueue_rollout(input={"id": 2}) # mode=None _r3 = await server.enqueue_rollout(input={"id": 3}, mode="test") # Sort by mode ascending (None values should be treated as empty string/0) async with session.get( f"{api_endpoint}/rollouts", params={"sort_by": "mode", "sort_order": "asc", "limit": -1} ) as resp: assert resp.status == 200 data = await resp.json() items = data["items"] assert len(items) == 3 # Items with None mode should come first (treated as 0) assert items[0]["rollout_id"] == r2.rollout_id assert items[0]["mode"] is None # Sort by mode descending async with session.get( f"{api_endpoint}/rollouts", params={"sort_by": "mode", "sort_order": "desc", "limit": -1} ) as resp: assert resp.status == 200 data = await resp.json() items = data["items"] assert len(items) == 3 # Items with actual mode should come first assert items[0]["mode"] is not None # Item with None should be last assert items[2]["rollout_id"] == r2.rollout_id assert items[2]["mode"] is None @pytest.mark.asyncio async def test_rollouts_sorting_by_unsupported_field( server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str], ) -> None: """Test sorting by a field that doesn't exist on the model.""" server, _client, session, api_endpoint = server_client # Create rollouts for i in range(3): await server.enqueue_rollout(input={"id": i}) # Try to sort by a non-existent field - should return 400 error async with session.get( f"{api_endpoint}/rollouts", params={"sort_by": "nonexistent_field", "sort_order": "asc", "limit": -1} ) as resp: assert resp.status == 400 data = await resp.json() assert "Invalid sort_by: nonexistent_field, allowed fields are: " in data["detail"] # Attempts Pagination and Sorting Tests @pytest.mark.asyncio async def test_attempts_pagination_basic( server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str], ) -> None: """Test basic pagination for attempts endpoint.""" server, _client, session, api_endpoint = server_client # Create rollout and multiple attempts rollout = await server.enqueue_rollout(input={"test": "data"}) for _ in range(5): await server.start_attempt(rollout.rollout_id) # Get first page async with session.get( f"{api_endpoint}/rollouts/{rollout.rollout_id}/attempts", params={"limit": 2, "offset": 0} ) as resp: assert resp.status == 200 data = await resp.json() assert data["total"] == 5 assert data["limit"] == 2 assert data["offset"] == 0 assert len(data["items"]) == 2 # Get all attempts async with session.get(f"{api_endpoint}/rollouts/{rollout.rollout_id}/attempts", params={"limit": -1}) as resp: assert resp.status == 200 data = await resp.json() assert data["total"] == 5 assert len(data["items"]) == 5 @pytest.mark.asyncio @pytest.mark.parametrize("method", ["get", "post"]) async def test_attempts_search_supports_get_and_post( server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str], method: str ) -> None: server, _client, session, api_endpoint = server_client rollout = await server.enqueue_rollout(input={"test": "attempt-search"}) await server.start_attempt(rollout.rollout_id) data = await _request_attempts_page(session, api_endpoint, rollout.rollout_id, method, {"limit": 1, "offset": 0}) assert data["total"] == 1 assert data["limit"] == 1 assert len(data["items"]) == 1 @pytest.mark.asyncio async def test_attempts_sorting( server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str], ) -> None: """Test sorting attempts by start_time.""" server, _client, session, api_endpoint = server_client # Create rollout and attempts rollout = await server.enqueue_rollout(input={"test": "data"}) attempts: List[AttemptedRollout] = [] for i in range(3): attempt = await server.start_attempt(rollout.rollout_id) attempts.append(attempt) await asyncio.sleep(0.01) # Sort by start_time descending (default) async with session.get( f"{api_endpoint}/rollouts/{rollout.rollout_id}/attempts", params={"sort_by": "start_time", "sort_order": "desc", "limit": -1}, ) as resp: assert resp.status == 200 data = await resp.json() items = data["items"] for i in range(len(items) - 1): assert items[i]["start_time"] >= items[i + 1]["start_time"] # Sort ascending async with session.get( f"{api_endpoint}/rollouts/{rollout.rollout_id}/attempts", params={"sort_by": "start_time", "sort_order": "asc", "limit": -1}, ) as resp: assert resp.status == 200 data = await resp.json() items = data["items"] for i in range(len(items) - 1): assert items[i]["start_time"] <= items[i + 1]["start_time"] @pytest.mark.asyncio async def test_attempts_sorting_with_none_values( server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str], ) -> None: """Test sorting attempts by fields that may have None values.""" server, _client, session, api_endpoint = server_client # Create rollout and attempts rollout = await server.enqueue_rollout(input={"test": "data"}) a1 = await server.start_attempt(rollout.rollout_id) a2 = await server.start_attempt(rollout.rollout_id) a3 = await server.start_attempt(rollout.rollout_id) # Set worker_id for some attempts, leave others with None await server.update_attempt(rollout_id=rollout.rollout_id, attempt_id=a1.attempt.attempt_id, worker_id="worker-1") await server.update_attempt(rollout_id=rollout.rollout_id, attempt_id=a3.attempt.attempt_id, worker_id="worker-2") # a2 remains with worker_id=None # Sort by worker_id ascending (None values should be treated as empty/0) async with session.get( f"{api_endpoint}/rollouts/{rollout.rollout_id}/attempts", params={"sort_by": "worker_id", "sort_order": "asc", "limit": -1}, ) as resp: assert resp.status == 200 data = await resp.json() items = data["items"] assert len(items) == 3 # Items with None worker_id should come first assert items[0]["attempt_id"] == a2.attempt.attempt_id assert items[0]["worker_id"] is None # Resources Pagination and Sorting Tests @pytest.mark.asyncio async def test_resources_pagination_basic( server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str], ) -> None: """Test basic pagination for resources endpoint.""" server, _client, session, api_endpoint = server_client # Create multiple resources for i in range(7): llm = LLM( resource_type="llm", endpoint="http://localhost:8080", model=f"model-v{i}", sampling_parameters={"temperature": 0.7}, ) await server.add_resources({"llm": llm}) # Get first page async with session.get(f"{api_endpoint}/resources", params={"limit": 3, "offset": 0}) as resp: assert resp.status == 200 data = await resp.json() assert data["total"] == 7 assert data["limit"] == 3 assert len(data["items"]) == 3 # Get all resources async with session.get(f"{api_endpoint}/resources", params={"limit": -1}) as resp: assert resp.status == 200 data = await resp.json() assert data["total"] == 7 assert len(data["items"]) == 7 @pytest.mark.asyncio async def test_resources_sorting_by_resources_id( server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str], ) -> None: """Test sorting resources by resources_id.""" server, _client, session, api_endpoint = server_client # Create resources for i in range(5): llm = LLM( resource_type="llm", endpoint="http://localhost:8080", model=f"model-{i}", sampling_parameters={"temperature": 0.7}, ) await server.add_resources({"llm": llm}) # Sort by resources_id ascending (default) async with session.get( f"{api_endpoint}/resources", params={"sort_by": "resources_id", "sort_order": "asc", "limit": -1} ) as resp: assert resp.status == 200 data = await resp.json() items = data["items"] for i in range(len(items) - 1): assert items[i]["resources_id"] <= items[i + 1]["resources_id"] # Sort descending async with session.get( f"{api_endpoint}/resources", params={"sort_by": "resources_id", "sort_order": "desc", "limit": -1} ) as resp: assert resp.status == 200 data = await resp.json() items = data["items"] for i in range(len(items) - 1): assert items[i]["resources_id"] >= items[i + 1]["resources_id"] @pytest.mark.asyncio async def test_resources_filter_by_resources_id_contains( server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str], ) -> None: """Test filtering resources by resources_id_contains.""" server, _client, session, api_endpoint = server_client # Create resources with specific IDs await server.update_resources("test-resource-001", {}) await server.update_resources("test-resource-002", {}) await server.update_resources("prod-resource-003", {}) # Filter by "test-" prefix async with session.get(f"{api_endpoint}/resources", params={"resources_id_contains": "test-", "limit": -1}) as resp: assert resp.status == 200 data = await resp.json() assert data["total"] == 2 ids = [item["resources_id"] for item in data["items"]] assert "test-resource-001" in ids assert "test-resource-002" in ids assert "prod-resource-003" not in ids # Filter by "-003" suffix async with session.get(f"{api_endpoint}/resources", params={"resources_id_contains": "-003", "limit": -1}) as resp: assert resp.status == 200 data = await resp.json() assert data["total"] == 1 assert data["items"][0]["resources_id"] == "prod-resource-003" @pytest.mark.asyncio async def test_resources_combined_filter_sort_and_pagination( server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str], ) -> None: """Test combining filter, sort, and pagination for resources.""" server, _client, session, api_endpoint = server_client # Create resources with pattern await server.update_resources("prod-app-001", {}) await asyncio.sleep(0.01) await server.update_resources("prod-app-002", {}) await asyncio.sleep(0.01) await server.update_resources("test-app-001", {}) await asyncio.sleep(0.01) await server.update_resources("prod-db-001", {}) # Filter by "prod-" and sort by resources_id, then paginate async with session.get( f"{api_endpoint}/resources", params={ "resources_id_contains": "prod-", "sort_by": "resources_id", "sort_order": "asc", "limit": 2, "offset": 0, }, ) as resp: assert resp.status == 200 data = await resp.json() assert data["total"] == 3 # 3 resources contain "prod-" assert data["limit"] == 2 assert data["offset"] == 0 assert len(data["items"]) == 2 # Should get first 2 when sorted by resources_id asc assert data["items"][0]["resources_id"] == "prod-app-001" assert data["items"][1]["resources_id"] == "prod-app-002" # Get the next page async with session.get( f"{api_endpoint}/resources", params={ "resources_id_contains": "prod-", "sort_by": "resources_id", "sort_order": "asc", "limit": 2, "offset": 2, }, ) as resp: assert resp.status == 200 data = await resp.json() assert data["total"] == 3 assert len(data["items"]) == 1 # Only 1 item remaining assert data["items"][0]["resources_id"] == "prod-db-001" # Spans Pagination, Sorting, and Filtering Tests @pytest.mark.asyncio async def test_spans_pagination_basic( server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str], ) -> None: """Test basic pagination for spans endpoint.""" server, _client, session, api_endpoint = server_client # Create rollout and attempt rollout = await server.start_rollout(input={"test": "data"}) # Add multiple spans for i in range(10): span = _make_span(rollout.rollout_id, rollout.attempt.attempt_id, i + 1, f"span-{i}") await server.add_span(span) # Get first page async with session.get( f"{api_endpoint}/spans", params={"rollout_id": rollout.rollout_id, "limit": 3, "offset": 0} ) as resp: assert resp.status == 200 data = await resp.json() assert data["total"] == 10 assert data["limit"] == 3 assert len(data["items"]) == 3 # Get all spans async with session.get(f"{api_endpoint}/spans", params={"rollout_id": rollout.rollout_id, "limit": -1}) as resp: assert resp.status == 200 data = await resp.json() assert data["total"] == 10 assert len(data["items"]) == 10 @pytest.mark.asyncio @pytest.mark.parametrize("method", ["get", "post"]) async def test_spans_search_supports_get_and_post( server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str], method: str ) -> None: server, _client, session, api_endpoint = server_client attempted = await server.start_rollout(input={"test": "span-search"}) attempt_id = attempted.attempt.attempt_id for seq in range(3): await server.add_span(_make_span(attempted.rollout_id, attempt_id, seq + 1, f"span-{seq}")) data = await _request_spans_page( session, api_endpoint, attempted.rollout_id, method, {"attempt_id": attempt_id, "limit": 2, "offset": 0}, ) assert data["total"] == 3 assert data["limit"] == 2 assert len(data["items"]) == 2 @pytest.mark.asyncio async def test_spans_sorting_by_start_time( server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str], ) -> None: """Test sorting spans by start_time.""" server, _client, session, api_endpoint = server_client # Create rollout and spans rollout = await server.start_rollout(input={"test": "data"}) for i in range(5): span = _make_span(rollout.rollout_id, rollout.attempt.attempt_id, i + 1, f"span-{i}") await server.add_span(span) # Sort descending (default) async with session.get( f"{api_endpoint}/spans", params={"rollout_id": rollout.rollout_id, "sort_by": "start_time", "sort_order": "desc", "limit": -1}, ) as resp: assert resp.status == 200 data = await resp.json() items = data["items"] # All our test spans have the same start_time, so just verify the structure assert len(items) == 5 @pytest.mark.asyncio async def test_spans_filter_by_trace_id( server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str], ) -> None: """Test filtering spans by trace_id.""" server, _client, session, api_endpoint = server_client # Create rollout and spans with different trace IDs rollout = await server.start_rollout(input={"test": "data"}) span1 = _make_span(rollout.rollout_id, rollout.attempt.attempt_id, 1, "span1") span1.trace_id = "trace-123" await server.add_span(span1) span2 = _make_span(rollout.rollout_id, rollout.attempt.attempt_id, 2, "span2") span2.trace_id = "trace-456" await server.add_span(span2) span3 = _make_span(rollout.rollout_id, rollout.attempt.attempt_id, 3, "span3") span3.trace_id = "trace-123" await server.add_span(span3) # Filter by exact trace_id async with session.get( f"{api_endpoint}/spans", params={"rollout_id": rollout.rollout_id, "trace_id": "trace-123", "limit": -1} ) as resp: assert resp.status == 200 data = await resp.json() assert data["total"] == 2 for item in data["items"]: assert item["trace_id"] == "trace-123" @pytest.mark.asyncio async def test_spans_filter_by_name_contains( server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str], ) -> None: """Test filtering spans by name_contains.""" server, _client, session, api_endpoint = server_client # Create rollout and spans with different names rollout = await server.start_rollout(input={"test": "data"}) span1 = _make_span(rollout.rollout_id, rollout.attempt.attempt_id, 1, "api_call_fetch") await server.add_span(span1) span2 = _make_span(rollout.rollout_id, rollout.attempt.attempt_id, 2, "database_query") await server.add_span(span2) span3 = _make_span(rollout.rollout_id, rollout.attempt.attempt_id, 3, "api_call_update") await server.add_span(span3) # Filter by name_contains async with session.get( f"{api_endpoint}/spans", params={"rollout_id": rollout.rollout_id, "name_contains": "api_call", "limit": -1} ) as resp: assert resp.status == 200 data = await resp.json() assert data["total"] == 2 for item in data["items"]: assert "api_call" in item["name"] @pytest.mark.asyncio async def test_spans_filter_logic_and( server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str], ) -> None: """Test spans filtering with AND logic.""" server, _client, session, api_endpoint = server_client # Create rollout and spans rollout = await server.start_rollout(input={"test": "data"}) span1 = _make_span(rollout.rollout_id, rollout.attempt.attempt_id, 1, "api_call") span1.trace_id = "trace-123" await server.add_span(span1) span2 = _make_span(rollout.rollout_id, rollout.attempt.attempt_id, 2, "api_call") span2.trace_id = "trace-456" await server.add_span(span2) span3 = _make_span(rollout.rollout_id, rollout.attempt.attempt_id, 3, "database") span3.trace_id = "trace-123" await server.add_span(span3) # Filter with AND logic: trace_id=trace-123 AND name contains "api" async with session.get( f"{api_endpoint}/spans", params={ "rollout_id": rollout.rollout_id, "trace_id": "trace-123", "name_contains": "api", "filter_logic": "and", "limit": -1, }, ) as resp: assert resp.status == 200 data = await resp.json() # Only span1 matches both conditions assert data["total"] == 1 assert data["items"][0]["name"] == "api_call" assert data["items"][0]["trace_id"] == "trace-123" @pytest.mark.asyncio async def test_spans_filter_logic_or( server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str], ) -> None: """Test spans filtering with OR logic.""" server, _client, session, api_endpoint = server_client # Create rollout and spans rollout = await server.start_rollout(input={"test": "data"}) span1 = _make_span(rollout.rollout_id, rollout.attempt.attempt_id, 1, "api_call") span1.trace_id = "trace-123" await server.add_span(span1) span2 = _make_span(rollout.rollout_id, rollout.attempt.attempt_id, 2, "api_call") span2.trace_id = "trace-456" await server.add_span(span2) span3 = _make_span(rollout.rollout_id, rollout.attempt.attempt_id, 3, "database") span3.trace_id = "trace-123" await server.add_span(span3) span4 = _make_span(rollout.rollout_id, rollout.attempt.attempt_id, 4, "other") span4.trace_id = "trace-789" await server.add_span(span4) # Filter with OR logic: trace_id=trace-123 OR name contains "api" async with session.get( f"{api_endpoint}/spans", params={ "rollout_id": rollout.rollout_id, "trace_id": "trace-123", "name_contains": "api", "filter_logic": "or", "limit": -1, }, ) as resp: assert resp.status == 200 data = await resp.json() # span1, span2, and span3 should match assert data["total"] == 3 @pytest.mark.asyncio async def test_spans_sorting_with_none_values( server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str], ) -> None: """Test sorting spans by fields that may have None values.""" server, _client, session, api_endpoint = server_client # Create rollout and spans rollout = await server.start_rollout(input={"test": "data"}) # Create spans with different parent_id values (parent_id can be None) span1 = _make_span(rollout.rollout_id, rollout.attempt.attempt_id, 1, "span1") span1.parent_id = None # Root span await server.add_span(span1) span2 = _make_span(rollout.rollout_id, rollout.attempt.attempt_id, 2, "span2") span2.parent_id = "parent-a" # Child span await server.add_span(span2) span3 = _make_span(rollout.rollout_id, rollout.attempt.attempt_id, 3, "span3") span3.parent_id = "parent-b" # Child span await server.add_span(span3) # Sort by parent_id ascending (None values should be treated as empty/0) async with session.get( f"{api_endpoint}/spans", params={"rollout_id": rollout.rollout_id, "sort_by": "parent_id", "sort_order": "asc", "limit": -1}, ) as resp: assert resp.status == 200 data = await resp.json() items = data["items"] assert len(items) == 3 # Items with None parent_id should come first (treated as 0) assert items[0]["sequence_id"] == 1 assert items[0]["parent_id"] is None # Sort by parent_id descending async with session.get( f"{api_endpoint}/spans", params={"rollout_id": rollout.rollout_id, "sort_by": "parent_id", "sort_order": "desc", "limit": -1}, ) as resp: assert resp.status == 200 data = await resp.json() items = data["items"] assert len(items) == 3 # Items with actual parent_id should come first assert items[0]["parent_id"] is not None # Item with None should be last assert items[2]["sequence_id"] == 1 assert items[2]["parent_id"] is None @pytest.mark.asyncio async def test_spans_sorting_by_unsupported_field( server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str], ) -> None: """Test sorting spans by a field that doesn't exist on the model.""" server, _client, session, api_endpoint = server_client # Create rollout and spans rollout = await server.start_rollout(input={"test": "data"}) for i in range(3): span = _make_span(rollout.rollout_id, rollout.attempt.attempt_id, i + 1, f"span-{i}") await server.add_span(span) # Try to sort by a non-existent field async with session.get( f"{api_endpoint}/spans", params={"rollout_id": rollout.rollout_id, "sort_by": "invalid_field", "sort_order": "asc", "limit": -1}, ) as resp: assert resp.status == 400 data = await resp.json() assert "Invalid sort_by: invalid_field, allowed fields are: " in data["detail"] # LightningStoreClient._request_json pagination metadata tests @pytest.mark.asyncio async def test_request_json_rollouts_returns_pagination_metadata( server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str], ) -> None: server, client, _session, _api_endpoint = server_client for idx in range(3): await server.enqueue_rollout(input={"idx": idx}) params = [ ("sort_by", "rollout_id"), ("sort_order", "asc"), ("limit", 1), ("offset", 1), ] data = await client._request_json("get", "/rollouts", params=params) # pyright: ignore[reportPrivateUsage] assert data["limit"] == 1 assert data["offset"] == 1 assert data["total"] == 3 assert len(data["items"]) == 1 @pytest.mark.asyncio async def test_request_json_attempts_returns_pagination_metadata( server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str], ) -> None: _server, client, _session, _api_endpoint = server_client attempted = await client.start_rollout(input={"payload": "attempts"}) await client.start_attempt(attempted.rollout_id) await client.start_attempt(attempted.rollout_id) params = [ ("sort_by", "sequence_id"), ("sort_order", "asc"), ("limit", 1), ("offset", 1), ] data = await client._request_json( # pyright: ignore[reportPrivateUsage] "get", f"/rollouts/{attempted.rollout_id}/attempts", params=params ) assert data["limit"] == 1 assert data["offset"] == 1 assert data["total"] == 3 assert len(data["items"]) == 1 @pytest.mark.asyncio async def test_request_json_resources_returns_pagination_metadata( server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str], ) -> None: _server, client, _session, _api_endpoint = server_client alpha = PromptTemplate(resource_type="prompt_template", template="alpha", engine="jinja") beta = PromptTemplate(resource_type="prompt_template", template="beta", engine="jinja") await client.update_resources("manual-alpha", {"prompt": alpha}) await client.update_resources("manual-beta", {"prompt": beta}) params = [ ("sort_by", "resources_id"), ("sort_order", "asc"), ("limit", 1), ("offset", 1), ] data = await client._request_json("get", "/resources", params=params) # pyright: ignore[reportPrivateUsage] assert data["limit"] == 1 assert data["offset"] == 1 assert data["total"] == 2 assert len(data["items"]) == 1 @pytest.mark.asyncio async def test_request_json_workers_returns_pagination_metadata( server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str], ) -> None: _server, client, _session, _api_endpoint = server_client for worker_id in ["worker-a", "worker-b", "worker-c"]: await client.update_worker(worker_id, heartbeat_stats={"cpu": 0.1}) params = [ ("sort_by", "worker_id"), ("sort_order", "asc"), ("limit", 1), ("offset", 1), ] data = await client._request_json("get", "/workers", params=params) # pyright: ignore[reportPrivateUsage] assert data["limit"] == 1 assert data["offset"] == 1 assert data["total"] == 3 assert len(data["items"]) == 1 @pytest.mark.asyncio async def test_request_json_spans_returns_pagination_metadata( server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str], ) -> None: server, client, _session, _api_endpoint = server_client attempted = await server.start_rollout(input={"span": "meta"}) attempt_id = attempted.attempt.attempt_id for idx in range(1, 4): await server.add_span(_make_span(attempted.rollout_id, attempt_id, idx, f"span-{idx}")) params = [ ("rollout_id", attempted.rollout_id), ("attempt_id", attempt_id), ("sort_by", "sequence_id"), ("sort_order", "asc"), ("limit", 1), ("offset", 1), ] data = await client._request_json("get", "/spans", params=params) # pyright: ignore[reportPrivateUsage] assert data["limit"] == 1 assert data["offset"] == 1 assert data["total"] == 3 assert len(data["items"]) == 1 # Update semantics tests @pytest.mark.asyncio async def test_update_rollout_distinguishes_unset_fields( server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str], ) -> None: server, _client, session, api_endpoint = server_client rollout = await server.enqueue_rollout(input={"id": "unset-check"}) async with session.post( f"{api_endpoint}/rollouts/{rollout.rollout_id}", json={"metadata": {"foo": "bar"}}, ) as resp: assert resp.status == 200 async with session.post( f"{api_endpoint}/rollouts/{rollout.rollout_id}", json={"status": None}, ) as resp: assert resp.status == 400 @pytest.mark.asyncio async def test_update_attempt_distinguishes_unset_fields( server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str], ) -> None: server, _client, session, api_endpoint = server_client attempted = await server.start_rollout(input={"id": "attempt-unset"}) attempt_id = attempted.attempt.attempt_id async with session.post( f"{api_endpoint}/rollouts/{attempted.rollout_id}/attempts/{attempt_id}", json={"status": "running"}, ) as resp: assert resp.status == 200 async with session.post( f"{api_endpoint}/rollouts/{attempted.rollout_id}/attempts/{attempt_id}", json={"worker_id": None}, ) as resp: assert resp.status == 400 # Client Compatibility Tests @pytest.mark.asyncio async def test_client_query_rollouts_extracts_items( server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str], ) -> None: """Test that client correctly extracts items from PaginatedResponse.""" server, client, _session, _api_endpoint = server_client # Create rollouts for i in range(5): await server.enqueue_rollout(input={"index": i}) # Query via client (should return PaginatedResult that behaves like a sequence) rollouts = await client.query_rollouts() assert isinstance(rollouts, PaginatedResult) assert rollouts.total == 5 assert len(rollouts) == 5 for rollout in rollouts: assert isinstance(rollout, Rollout) @pytest.mark.asyncio async def test_client_query_with_filters( server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str], ) -> None: """Test that client correctly passes filters to server.""" server, client, _session, _api_endpoint = server_client # Create rollouts with different statuses r1 = await server.enqueue_rollout(input={"id": 1}) r2 = await server.enqueue_rollout(input={"id": 2}) await server.update_rollout(rollout_id=r1.rollout_id, status="succeeded") # Query via client with status filter rollouts = await client.query_rollouts(status=["succeeded"]) assert len(rollouts) == 1 assert rollouts[0].rollout_id == r1.rollout_id # Query via client with rollout_ids filter rollouts = await client.query_rollouts(rollout_ids=[r2.rollout_id]) assert len(rollouts) == 1 assert rollouts[0].rollout_id == r2.rollout_id @pytest.mark.asyncio @pytest.mark.parametrize("method", ["get", "post"]) async def test_workers_search_supports_get_and_post( server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str], method: str ) -> None: server, _client, session, api_endpoint = server_client await server.update_worker("worker-1", heartbeat_stats={"cpu": 0.5}) await server.update_worker("worker-2", heartbeat_stats={"cpu": 0.7}) data = await _request_workers_page( session, api_endpoint, method, {"limit": 1, "offset": 0, "sort_by": "worker_id", "sort_order": "asc"}, ) assert data["total"] == 2 assert data["limit"] == 1 assert len(data["items"]) == 1 @pytest.mark.asyncio async def test_workers_endpoint_supports_updates( server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str], ) -> None: _server, _client, session, api_endpoint = server_client async with session.post( f"{api_endpoint}/workers/worker-1", json={"heartbeat_stats": {"cpu": 0.7}}, ) as resp: assert resp.status == 200 created = await resp.json() assert created["worker_id"] == "worker-1" assert created["status"] == "unknown" assert created["heartbeat_stats"] == {"cpu": 0.7} first_heartbeat = created["last_heartbeat_time"] async with session.get(f"{api_endpoint}/workers") as resp: assert resp.status == 200 data = await resp.json() workers = data["items"] assert len(workers) == 1 assert workers[0]["worker_id"] == "worker-1" async with session.post( f"{api_endpoint}/workers/worker-1", json={"heartbeat_stats": {"cpu": 0.8}}, ) as resp: assert resp.status == 200 updated = await resp.json() assert updated["last_heartbeat_time"] >= first_heartbeat async with session.get(f"{api_endpoint}/workers") as resp: assert resp.status == 200 data = await resp.json() workers = data["items"] assert workers[0]["status"] == "unknown" @pytest.mark.asyncio async def test_workers_endpoint_rejects_none_stats( server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str], ) -> None: _server, _client, session, api_endpoint = server_client async with session.post( f"{api_endpoint}/workers/worker-err", json={"heartbeat_stats": None}, ) as resp: assert resp.status == 400 @pytest.mark.asyncio async def test_get_worker_by_id_restful( server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str], ) -> None: server, _client, session, api_endpoint = server_client await server.update_worker("worker-fetch", heartbeat_stats={"cpu": 0.4}) async with session.get(f"{api_endpoint}/workers/worker-fetch") as resp: assert resp.status == 200 data = await resp.json() assert data["worker_id"] == "worker-fetch" async with session.get(f"{api_endpoint}/workers/missing") as resp: assert resp.status == 200 data = await resp.json() assert data is None @pytest.mark.asyncio async def test_workers_endpoint_filter_and_sort( server_client: Tuple[LightningStoreServer, LightningStoreClient, aiohttp.ClientSession, str], ) -> None: server, _client, session, api_endpoint = server_client # Worker A: finishes an attempt and becomes idle. await server.update_worker("worker-a", heartbeat_stats={"cpu": 0.1}) await server.enqueue_rollout(input={"worker": "a"}) claimed_a = await server.dequeue_rollout(worker_id="worker-a") assert claimed_a is not None await server.update_attempt( claimed_a.rollout_id, claimed_a.attempt.attempt_id, worker_id="worker-a", status="succeeded" ) # Worker B: currently busy on an attempt. await server.update_worker("worker-b", heartbeat_stats={"cpu": 0.9}) await server.enqueue_rollout(input={"worker": "b"}) claimed_b = await server.dequeue_rollout(worker_id="worker-b") assert claimed_b is not None await server.update_attempt(claimed_b.rollout_id, claimed_b.attempt.attempt_id, worker_id="worker-b") # Worker C: also busy. await server.update_worker("worker-c", heartbeat_stats={"cpu": 0.2}) await server.enqueue_rollout(input={"worker": "c"}) claimed_c = await server.dequeue_rollout(worker_id="worker-c") assert claimed_c is not None await server.update_attempt(claimed_c.rollout_id, claimed_c.attempt.attempt_id, worker_id="worker-c") async with session.get( f"{api_endpoint}/workers", params={"status_in": ["busy"], "worker_id_contains": "worker", "sort_by": "worker_id", "sort_order": "desc"}, ) as resp: assert resp.status == 200 data = await resp.json() items = data["items"] assert [w["worker_id"] for w in items] == ["worker-c", "worker-b"] async with session.get( f"{api_endpoint}/workers", params={"limit": 1, "offset": 1, "sort_by": "worker_id", "sort_order": "asc"}, ) as resp: assert resp.status == 200 data = await resp.json() assert data["limit"] == 1 assert data["offset"] == 1 assert len(data["items"]) == 1