xc-model-auto-adaptation-agent/test_main.py

import io
import json
import unittest
from contextlib import redirect_stdout
from unittest.mock import patch

import main as agent_main
from main import (
    AGENT_NAME,
    PlatformAPIError,
    _log_event,
    _safe_model_name,
    build_adaptation_plan,
    resolve_runtime_token,
    scan_once,
    select_candidate,
)


class AutoAdaptationTests(unittest.TestCase):
    @staticmethod
    def candidate_payload():
        return {
            "code": 0,
            "data": [
                {
                    "modelId": "Qwen/Qwen3-30B-Instruct",
                    "taskLevelsInfo": {"taskLevelChineseName": "文本生成"},
                },
                {
                    "modelId": "Qwen/Qwen3-4B-Instruct-2507",
                    "taskLevelsInfo": {"taskLevelChineseName": "文本生成"},
                },
                {
                    "modelId": "example/image-model",
                    "taskLevelsInfo": {"taskLevelChineseName": "文本生成图片"},
                },
            ],
        }

    def test_official_runtime_token_name_is_supported(self):
        token = resolve_runtime_token({"EXTERNAL_SERVICE_TOKEN": "private-value"})
        self.assertEqual(token, "private-value")

    def test_runtime_token_file_is_supported(self):
        token = resolve_runtime_token(
            {"XC_TOKEN_FILE": "/run/secrets/xc-token"},
            read_text=lambda path: "mounted-private-value\n"
            if path == "/run/secrets/xc-token"
            else "",
        )
        self.assertEqual(token, "mounted-private-value")

    def test_placeholder_runtime_token_is_rejected(self):
        token = resolve_runtime_token({"EXTERNAL_SERVICE_TOKEN": "tmp"})
        self.assertEqual(token, "")

    def test_explicit_modelhub_token_takes_precedence(self):
        token = resolve_runtime_token(
            {
                "MODELHUB_XC_TOKEN": "preferred",
                "EXTERNAL_SERVICE_TOKEN": "fallback",
            }
        )
        self.assertEqual(token, "preferred")

    def test_candidate_selection_prefers_small_supported_text_model(self):
        candidate = select_candidate(self.candidate_payload())
        self.assertIsNotNone(candidate)
        self.assertEqual(candidate["model_id"], "Qwen/Qwen3-4B-Instruct-2507")

    def test_scan_blocks_submission_without_runtime_token(self):
        calls = []

        def fake_request(method, url, **kwargs):
            calls.append((method, url, kwargs))
            return self.candidate_payload()

        with patch.object(agent_main, "XC_TOKEN", ""):
            result = scan_once(fake_request)
        self.assertEqual(result["action"], "submission_blocked")
        self.assertIn("runtime_token_missing", result["reasons"])
        self.assertEqual([call[0] for call in calls], ["GET"])

    def test_scan_deduplicates_before_submission(self):
        calls = []

        def fake_request(method, url, **kwargs):
            calls.append((method, url, kwargs))
            if "top/models" in url:
                return self.candidate_payload()
            return {"code": 0, "data": {"records": [{"id": 1}]}}

        with (
            patch.object(agent_main, "XC_TOKEN", "private-token"),
            patch.object(agent_main, "STRATEGY_ID", "strategy-id"),
            patch.object(agent_main, "TARGET_GPU", "verified-gpu"),
            patch.object(agent_main, "CONFIG_PARAMS", "framework: vllm\nsut_config: test"),
            patch.object(agent_main, "AUTO_SUBMIT_ENABLED", True),
        ):
            result = scan_once(fake_request)
        self.assertEqual(result["action"], "duplicate_skipped")
        self.assertEqual([call[0] for call in calls], ["GET", "GET"])

    def test_scan_submits_exactly_one_task(self):
        calls = []

        def fake_request(method, url, **kwargs):
            calls.append((method, url, kwargs))
            if "top/models" in url:
                return self.candidate_payload()
            if method == "GET":
                return {"code": 0, "data": {"records": []}}
            return {"code": 0, "data": {"id": 987, "status": "created"}}

        with (
            patch.object(agent_main, "XC_TOKEN", "private-token"),
            patch.object(agent_main, "STRATEGY_ID", "strategy-id"),
            patch.object(agent_main, "TARGET_GPU", "verified-gpu"),
            patch.object(agent_main, "CONFIG_PARAMS", "framework: vllm\nsut_config: test"),
            patch.object(agent_main, "AUTO_SUBMIT_ENABLED", True),
        ):
            result = scan_once(fake_request)
        self.assertEqual(result, {
            "action": "submitted",
            "model": "Qwen/Qwen3-4B-Instruct-2507",
            "task_id": 987,
        })
        self.assertEqual([call[0] for call in calls], ["GET", "GET", "POST"])
        submitted = calls[-1][2]["payload"]
        self.assertEqual(submitted["strategyId"], "strategy-id")
        self.assertEqual(submitted["targetGpu"], "verified-gpu")
        self.assertIn("sut_config", submitted["configParams"])

    def test_platform_business_error_is_not_reported_as_success(self):
        def fake_request(method, url, **kwargs):
            if "top/models" in url:
                return self.candidate_payload()
            return {"code": 403, "message": "forbidden"}

        with (
            patch.object(agent_main, "XC_TOKEN", "private-token"),
            patch.object(agent_main, "STRATEGY_ID", "strategy-id"),
            patch.object(agent_main, "TARGET_GPU", "verified-gpu"),
            patch.object(agent_main, "CONFIG_PARAMS", "framework: vllm"),
            patch.object(agent_main, "AUTO_SUBMIT_ENABLED", True),
        ):
            with self.assertRaises(PlatformAPIError):
                scan_once(fake_request)

    def test_adaptation_plan_is_execution_oriented(self):
        with (
            patch.object(agent_main, "XC_TOKEN", "private-token"),
            patch.object(agent_main, "STRATEGY_ID", "strategy-id"),
        ):
            plan = build_adaptation_plan({"model_name": "Qwen/Qwen3-4B"})
        self.assertEqual(plan["task_type"], "text-generation")
        self.assertIn("提交平台验证任务", plan["execution"])
        self.assertTrue(plan["automatic_submission_ready"])

    def test_model_identifier_removes_credentials_and_query(self):
        value = "https://user:pass@example.com/org/model?token=secret#fragment"
        self.assertEqual(_safe_model_name(value), "https://example.com/org/model")

    def test_structured_log_does_not_need_secret_values(self):
        stream = io.StringIO()
        with redirect_stdout(stream):
            _log_event("scanner_started", runtime_token_present=True)
        record = json.loads(stream.getvalue())
        self.assertEqual(record["agent"], AGENT_NAME)
        self.assertTrue(record["runtime_token_present"])
        self.assertNotIn("token", record)


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