Mujoco_WASM/decision_server/tests/test_hosted_commands.py

import asyncio
import tempfile
import unittest
from pathlib import Path
from unittest.mock import AsyncMock, patch

from aiohttp.test_utils import TestClient, TestServer

from decision_server.commands import checked
from decision_server.connections import Connection
from decision_server.hosted_budget import HostedBudget
from decision_server.protocol import DecisionError
from decision_server.server import PREFIX
from decision_server.tests import test_website
from decision_server.web_server import MANAGER, create_website_app


class HostedTests(unittest.IsolatedAsyncioTestCase):
    visitor = test_website.WebsiteTests.visitor

    async def asyncSetUp(self):
        self.temp = tempfile.TemporaryDirectory()
        self.path = Path(self.temp.name) / "budget.sqlite"
        self.defaults = {
            "llm": Connection(
                "responses", "https://api.deepseek.com", "deepseek-flash", "llm-test-secret"
            ),
            "jev": Connection(
                "openrouter-decisions",
                "https://openrouter.ai/api/alpha/decisions",
                "typesafe/jev-1.13",
                "jev-test-secret",
            ),
        }
        self.origin = "https://site.test"
        self.headers = {"Host": "site.test", "Origin": self.origin}
        self.app = create_website_app(
            self.origin, defaults=self.defaults, budget=HostedBudget(self.path, 4, 8)
        )
        self.manager = self.app[MANAGER]
        self.client = TestClient(TestServer(self.app))
        await self.client.start_server()

    async def asyncTearDown(self):
        await self.client.close()
        self.temp.cleanup()

    def request(self, ident="request1"):
        return {
            "instruction": "前进0.5m",
            "stamp": {
                "runId": "run1",
                "requestId": ident,
                "sceneRevision": 0,
                "sequence": 0,
                "planRevision": 0,
            },
        }

    async def test_defaults_locked_and_private(self):
        a, av = await self.visitor()
        b, bv = await self.visitor()
        response = await self.client.get(PREFIX + "/status", headers=a)
        text = await response.text()
        self.assertTrue((await response.json())["ready"])
        for conn in self.defaults.values():
            self.assertNotIn(conn.key, text)
        for method, path in [
            ("put", "/configuration"),
            ("post", "/test"),
            ("post", "/codex/login"),
            ("get", "/models"),
        ]:
            response = await getattr(self.client, method)(PREFIX + path, headers=a, json={})
            self.assertEqual(response.status, 403)
        await self.client.delete(PREFIX + "/session", headers=a)
        self.assertFalse(av.service.connections.values)
        self.assertTrue(bv.service.connections.values)
        self.assertEqual(self.manager.defaults, self.defaults)

    async def test_real_shape_dual_calls_no_retry_and_shared_persistent_budget(self):
        llm = AsyncMock(
            return_value=(
                {"action": "move", "value": 0.5, "summary": "前进0.5米"},
                {"input_tokens": 1},
            )
        )
        gate = AsyncMock(return_value={"answers": {"gate": {"choice": "allow"}}, "usage": {}})
        with (
            patch("decision_server.providers.openai.structured", llm),
            patch("decision_server.providers.jev.post", gate),
        ):
            for i in range(3):
                headers, _ = await self.visitor()
                response = await self.client.post(
                    PREFIX + "/command", json=self.request(str(i)), headers=headers
                )
                self.assertEqual(response.status, 200 if i < 2 else 429)
            self.assertEqual(llm.await_count, 2)
            self.assertEqual(gate.await_count, 2)
        with self.assertRaisesRegex(DecisionError, "shared_budget_exceeded"):
            HostedBudget(self.path, 4, 8).reserve()

    async def test_cancel_owns_jev_and_other_session_cannot_cancel(self):
        a, _ = await self.visitor()
        b, bv = await self.visitor()
        started = asyncio.Event()

        async def delayed(*_):
            started.set()
            await asyncio.Event().wait()

        with (
            patch(
                "decision_server.providers.openai.structured",
                AsyncMock(return_value=({"action": "move", "value": 0.5, "summary": "前进"}, {})),
            ),
            patch("decision_server.providers.jev.post", side_effect=delayed),
        ):
            job = asyncio.create_task(
                self.client.post(PREFIX + "/command", headers=b, json=self.request())
            )
            await asyncio.wait_for(started.wait(), 2)
            identity = {"runId": "run1", "requestId": "request1"}
            response = await self.client.post(PREFIX + "/cancel", headers=a, json=identity)
            self.assertFalse((await response.json())["cancelled"])
            response = await self.client.post(PREFIX + "/cancel", headers=b, json=identity)
            self.assertTrue((await response.json())["cancelled"])
            self.assertEqual((await job).status, 409)
            self.assertFalse(bv.service.active)

    def test_output_rejects_bad_direction_units_values_and_ambiguity(self):
        for value in [True, float("nan"), 2, -0.5]:
            with self.assertRaises(DecisionError):
                checked({"action": "move", "value": value, "summary": "走"}, "前进0.5m")
        self.assertEqual(
            checked({"action": "turn", "value": 90, "summary": "转"}, "转动90度")["action"],
            "clarify",
        )
        self.assertEqual(
            checked({"action": "move", "value": 0.5, "summary": "走"}, "不要前进0.5米")["action"],
            "clarify",
        )
        self.assertEqual(
            checked({"action": "move", "value": 0.5, "summary": "走"}, "前进50cm")["value"], 0.5
        )
        with self.assertRaises(DecisionError):
            checked({"action": "turn", "value": 90, "summary": "转"}, "右转90度")