xc-tianga100-advisor-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,
    CARD_MODEL,
    CARD_VENDOR,
    ValidationError,
    _log_event,
    _safe_model_name,
    analyze,
    scan_once,
    select_candidate,
)


class AnalyzeTests(unittest.TestCase):
    def complete_payload(self):
        return {
            "model_name": "example/model",
            "hardware": f"{CARD_VENDOR} {CARD_MODEL}",
            "framework": "PyTorch",
            "backend": "vendor-backend",
            "sdk_version": "provided-by-user",
            "driver_version": "provided-by-user",
        }

    def test_complete_preflight_targets_this_card(self):
        result = analyze(self.complete_payload())
        self.assertEqual(result["verdict"], "preflight_ready")
        self.assertTrue(result["target_matches"])
        self.assertEqual(result["official_target"]["model"], CARD_MODEL)

    def test_mismatched_target_is_rejected(self):
        payload = self.complete_payload()
        payload["hardware"] = "其他厂商 其他卡型"
        result = analyze(payload)
        self.assertEqual(result["verdict"], "target_mismatch")
        self.assertFalse(result["target_matches"])

    def test_missing_information_is_explicit(self):
        result = analyze({})
        self.assertEqual(result["verdict"], "information_required")
        self.assertIn("sdk_version", result["missing_fields"])
        self.assertIn("driver_version", result["missing_fields"])

    def test_multicard_and_quantization_risks_are_reported(self):
        payload = self.complete_payload()
        payload.update({"cards": 2, "precision": "int8"})
        result = analyze(payload)
        joined = " ".join(result["risks"])
        self.assertIn("多卡", joined)
        self.assertIn("量化", joined)

    def test_invalid_card_count_is_rejected(self):
        with self.assertRaises(ValidationError):
            analyze({"cards": 0})

    def test_model_log_identifier_removes_url_credentials(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_is_valid_json(self):
        stream = io.StringIO()
        with redirect_stdout(stream):
            _log_event("analysis_completed", model="example/model", verdict="ready")
        record = json.loads(stream.getvalue())
        self.assertEqual(record["agent"], AGENT_NAME)
        self.assertEqual(record["event"], "analysis_completed")
        self.assertEqual(record["model"], "example/model")

    @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_candidate_selection_prefers_smaller_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_skips_write_without_private_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_skipped")
        self.assertIn("xc_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": [{"taskId": "existing"}]}}

        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),
        ):
            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_after_empty_history(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": {"taskId": "new-task"}}

        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),
        ):
            result = scan_once(fake_request)
        self.assertEqual(result["action"], "submitted")
        self.assertEqual(result["task_id"], "new-task")
        self.assertEqual([call[0] for call in calls], ["GET", "GET", "POST"])
        submitted = calls[-1][2]["payload"]
        self.assertEqual(submitted["targetGpu"], "verified-gpu")
        self.assertEqual(submitted["strategyId"], "strategy-id")


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