fengqun/tests/test_newapi_agnet.py

import unittest

from swarm_minimal.core import InMemorySwarmStore, SwarmCoordinator, Task, default_agents
from swarm_minimal.newapi_agnet import (
    NewApiAgnet,
    NewApiChannelConfig,
    build_model_test_agnets,
    discover_newapi_models,
    select_distinct_models,
)


class FakeHttpClient:
    def __init__(self) -> None:
        self.calls = []
        self.model_response = {
            "data": [
                {"id": "model-alpha"},
                {"id": "model-beta"},
                {"id": "model-gamma"},
            ]
        }

    def get_json(self, url, headers, timeout):
        self.calls.append(
            {
                "method": "GET",
                "url": url,
                "headers": headers,
                "timeout": timeout,
            }
        )
        return self.model_response

    def post_json(self, url, headers, payload, timeout):
        self.calls.append(
            {
                "method": "POST",
                "url": url,
                "headers": headers,
                "payload": payload,
                "timeout": timeout,
            }
        )
        return {
            "choices": [
                {
                    "message": {
                        "content": "Mocked NewAPI Agnet verification passed.",
                    }
                }
            ]
        }


class NewApiAgnetTest(unittest.TestCase):
    def test_newapi_agnet_posts_openai_compatible_chat_request(self) -> None:
        fake_http = FakeHttpClient()
        config = NewApiChannelConfig(
            base_url="http://newapi.example.test",
            api_key="test-secret",
            model="test-model",
            timeout_seconds=3,
        )
        agnet = NewApiAgnet(config, http_client=fake_http)

        output, score = agnet.run_task(
            task=type("TaskLike", (), {"kind": "verify", "input": "minimal swarm"})(),
            shared_state={"run:1:status": "running"},
        )

        self.assertEqual(output, "Mocked NewAPI Agnet verification passed.")
        self.assertEqual(score, 0.9)
        self.assertEqual(len(fake_http.calls), 1)
        call = fake_http.calls[0]
        self.assertEqual(call["method"], "POST")
        self.assertEqual(call["url"], "http://newapi.example.test/v1/chat/completions")
        self.assertEqual(call["headers"]["Authorization"], "Bearer test-secret")
        self.assertEqual(call["payload"]["model"], "test-model")
        self.assertIn("messages", call["payload"])
        self.assertEqual(call["timeout"], 3)

    def test_newapi_agnet_can_replace_verifier_in_minimal_swarm(self) -> None:
        fake_http = FakeHttpClient()
        config = NewApiChannelConfig(
            base_url="http://newapi.example.test",
            api_key="test-secret",
            model="test-model",
        )
        store = InMemorySwarmStore()
        agents = [
            default_agents()[0],
            default_agents()[1],
            NewApiAgnet(config, agent_id="newapi-verifier", capability="verify", http_client=fake_http).as_agent(),
        ]

        run_id = SwarmCoordinator(store=store, agents=agents).submit_goal("test agnet")
        result = SwarmCoordinator(store=store, agents=agents).run_until_converged(run_id)

        self.assertEqual(result.completed_tasks, 3)
        self.assertEqual(result.accepted_output, "Mocked NewAPI Agnet verification passed.")
        self.assertEqual(len(fake_http.calls), 1)

    def test_config_summary_redacts_newapi_key(self) -> None:
        summary = NewApiChannelConfig(
            base_url="http://newapi.example.test",
            api_key="test-secret",
        ).redacted_summary()

        self.assertEqual(summary["api_key"], "<redacted>")
        self.assertNotIn("test-secret", str(summary))

    def test_discovers_models_from_openai_compatible_endpoint(self) -> None:
        fake_http = FakeHttpClient()
        config = NewApiChannelConfig(
            base_url="http://newapi.example.test",
            api_key="test-secret",
            timeout_seconds=5,
        )

        models = discover_newapi_models(config, http_client=fake_http)

        self.assertEqual(models, ["model-alpha", "model-beta", "model-gamma"])
        self.assertEqual(fake_http.calls[0]["method"], "GET")
        self.assertEqual(fake_http.calls[0]["url"], "http://newapi.example.test/v1/models")
        self.assertEqual(fake_http.calls[0]["headers"]["Authorization"], "Bearer test-secret")
        self.assertEqual(fake_http.calls[0]["timeout"], 5)

    def test_selects_three_distinct_models(self) -> None:
        selected = select_distinct_models(["a", "b", "a", "c", "d"], count=3)

        self.assertEqual(selected, ["a", "b", "c"])

    def test_three_newapi_agnets_use_three_different_models(self) -> None:
        fake_http = FakeHttpClient()
        config = NewApiChannelConfig(
            base_url="http://newapi.example.test",
            api_key="test-secret",
            model="default-model",
        )
        models = discover_newapi_models(config, http_client=fake_http)
        selected_models = select_distinct_models(models, count=3)
        agents = build_model_test_agnets(config, models=selected_models, http_client=fake_http)
        store = InMemorySwarmStore()
        coordinator = SwarmCoordinator(store=store, agents=agents)
        run_id = coordinator.submit_goal("test three models")
        for index, model in enumerate(selected_models):
            store.add_task(Task(kind=f"model_test_{index + 1}", input=f"model={model}"))

        result = coordinator.run_until_converged(run_id)

        post_calls = [call for call in fake_http.calls if call["method"] == "POST"]
        self.assertEqual(len(post_calls), 3)
        self.assertEqual([call["payload"]["model"] for call in post_calls], selected_models)
        self.assertEqual(result.completed_tasks, 3)
        self.assertEqual(result.accepted_output, "Mocked NewAPI Agnet verification passed.")


if __name__ == "__main__":
    unittest.main()