# Copyright 2026 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # https://www.apache.org/licenses/LICENSE-2.0 from __future__ import annotations import unittest from collections import defaultdict from collections.abc import Mapping from typing import Any from artemis_client import ( ArtemisClient, NotFoundError, ProtocolError, TaskRejectedError, TaskTimeoutError, ) class FakeTransport: def __init__(self) -> None: self.responses: dict[tuple[str, str], list[Any]] = defaultdict(list) self.calls: list[tuple[str, str, Mapping[str, Any] | None]] = [] def add(self, method: str, path: str, *responses: Any) -> None: self.responses[(method, path)].extend(responses) def request( self, method: str, path: str, *, json_body: Mapping[str, Any] | None = None, ) -> Any: self.calls.append((method, path, json_body)) queued = self.responses[(method, path)] if not queued: raise AssertionError(f"Unexpected request: {method} {path}") response = queued.pop(0) if isinstance(response, BaseException): raise response return response class ArtemisClientTests(unittest.IsolatedAsyncioTestCase): def setUp(self) -> None: self.transport = FakeTransport() self.client = ArtemisClient( "https://artemis.example.test", poll_interval=0.001, transport=self.transport, ) async def test_submit_sends_idempotent_session_id(self) -> None: task_id = "00000000-0000-4000-8000-000000000123" self.transport.add( "POST", "/api/run", { "status": "started", "tasks": [ { "session_id": task_id, "status": "pending", "device_serial": "pixel-8", } ], }, ) handle = await self.client.submit( "Open Settings", task_id=task_id, device_serial="pixel-8", options={"record_video": True}, ) self.assertEqual(handle.task_id, task_id) self.assertEqual(handle.device_serial, "pixel-8") body = self.transport.calls[0][2] assert body is not None self.assertEqual(body["session_id"], task_id) self.assertEqual(body["ingress"], "python_sdk") self.assertEqual(body["options"], {"record_video": True}) async def test_submit_rejected_task_raises_specific_error(self) -> None: self.transport.add( "POST", "/api/run", { "status": "rejected", "error": "Device is offline", "tasks": [], }, ) with self.assertRaisesRegex(TaskRejectedError, "Device is offline"): await self.client.submit("Open Settings") async def test_submit_rejects_non_uuid_task_id(self) -> None: with self.assertRaisesRegex(ValueError, "valid UUID"): await self.client.submit("Open Settings", task_id="not-a-uuid") async def test_run_finds_queued_task_then_reads_terminal_session(self) -> None: task_id = "00000000-0000-4000-8000-000000000124" self.transport.add( "POST", "/api/run", { "status": "started", "tasks": [{"session_id": task_id, "status": "pending"}], }, ) self.transport.add( "GET", f"/api/sessions/{task_id}", NotFoundError(404, "not created yet"), { "session_id": task_id, "status": "completed", "goal": "Open Settings", "current_turn": 4, "summary": "Battery page opened", }, ) self.transport.add( "GET", "/api/status", { "status": "running", "queue": [{"session_id": task_id, "status": "pending"}], }, ) result = await self.client.run("Open Settings", task_id=task_id, timeout=1) self.assertTrue(result.done) self.assertTrue(result.succeeded) self.assertEqual(result.turns, 4) self.assertEqual(result.output, "Battery page opened") async def test_get_task_returns_launching_when_not_visible_yet(self) -> None: self.transport.add( "GET", "/api/sessions/new-task", NotFoundError(404, "missing"), ) self.transport.add("GET", "/api/status", {"status": "idle", "queue": []}) result = await self.client.get_task("new-task") self.assertEqual(result.status, "launching") self.assertFalse(result.done) async def test_wait_for_task_times_out(self) -> None: class LaunchingClient(ArtemisClient): async def get_task(self, task_id: str): # type: ignore[override] from artemis_client import TaskResult return TaskResult(task_id=task_id, status="running") client = LaunchingClient( "https://artemis.example.test", poll_interval=0.001, transport=self.transport, ) with self.assertRaises(TaskTimeoutError): await client.wait_for_task("slow-task", timeout=0.003) async def test_list_devices_accepts_legacy_shape(self) -> None: self.transport.add( "GET", "/api/devices", { "devices": [ { "serial": "emulator-5554", "state": "device", "model": "Pixel_8", "busy": True, } ] }, ) devices = await self.client.list_devices() self.assertEqual(len(devices), 1) self.assertEqual(devices[0].serial, "emulator-5554") self.assertTrue(devices[0].busy) async def test_capabilities_falls_back_for_legacy_server(self) -> None: self.transport.add( "GET", "/api/v1/capabilities", NotFoundError(404, "not implemented"), ) capabilities = await self.client.capabilities() self.assertEqual(capabilities.api_version, "legacy") self.assertTrue(capabilities.supports("tasks.submit")) async def test_health_uses_fast_scheduler_endpoint(self) -> None: self.transport.add("GET", "/api/status", {"status": "idle"}) health = await self.client.health() self.assertEqual(health["status"], "idle") async def test_readiness_uses_full_diagnostics_endpoint(self) -> None: self.transport.add( "GET", "/api/system/readiness", {"overall_status": "ready"}, ) readiness = await self.client.readiness() self.assertEqual(readiness["overall_status"], "ready") async def test_invalid_devices_payload_is_rejected(self) -> None: self.transport.add("GET", "/api/devices", {"devices": "not-a-list"}) with self.assertRaises(ProtocolError): await self.client.list_devices() async def test_stop_targets_one_session(self) -> None: self.transport.add("POST", "/api/stop", {"status": "stopped"}) stopped = await self.client.stop("task-stop") self.assertTrue(stopped) self.assertEqual( self.transport.calls[0], ("POST", "/api/stop", {"session_id": "task-stop"}), ) if __name__ == "__main__": unittest.main()