xc_agent_platform_demo/tests/test_reject_marks.py

import io
import json
import os
from pathlib import Path
import subprocess
import sys
import tempfile
import unittest
import urllib.error
from unittest.mock import Mock, patch

from src.config_params import GPU_ENUMS
from src.live_cycle import run_cycle, _submitted, _marked
from src.marks import load_marks, PERMANENT_PAIR_CODES
from src.precheck import precheck_model


class TestMarks(unittest.TestCase):
    def test_baked_marks(self):
        marks = load_marks()
        self.assertEqual(len(marks), 8)
        self.assertTrue(all(g in GPU_ENUMS.values() for m, g in marks))
        self.assertEqual(marks['madeagents/hammer2.1-0.5b', 'Biren_166m'], 60027)
        self.assertEqual(PERMANENT_PAIR_CODES, {60012, *range(60019, 60030)})

    def test_invalid_marks(self):
        with tempfile.TemporaryDirectory() as directory:
            path = Path(directory) / 'marks.json'
            with self.assertRaises(ValueError):
                load_marks(path)
            valid = dict(model='org/model', gpu='Biren_166m', code=60027)
            for data in ('bad json', [], {}, {'version': 1, 'marks': {}},
                         {'version': 1, 'marks': [None]},
                         *({'version': 1, 'marks': [dict(valid, **update)]} for update in
                           ({'model': ''}, {'model': 4}, {'gpu': 'unknown'}, {'gpu': []}, {'code': True}, {'code': '60027'}))):
                path.write_text(data if isinstance(data, str) else json.dumps(data))
                with self.subTest(data=data), self.assertRaises(ValueError):
                    load_marks(path)


class TestPrecheck(unittest.TestCase):
    def check(self, body, raw=False):
        response = Mock(status=200)
        response.read.return_value = body if raw else json.dumps(body).encode()
        response.__enter__ = Mock(return_value=response)
        response.__exit__ = Mock(return_value=False)
        opener = Mock(return_value=response)
        result = precheck_model('org/model name', token='synthetic-secret', opener=opener)
        request = opener.call_args.args[0]
        self.assertEqual(request.full_url, 'https://modelhub.org.cn/api/computility/models/search-by-model-id?modelId=org%2Fmodel%20name')
        self.assertEqual(request.get_header('Xc-token'), 'synthetic-secret')
        self.assertEqual(request.get_header('Accept'), 'application/json')
        self.assertEqual(request.method, 'GET')
        self.assertEqual(opener.call_args.kwargs, {'timeout': 30})
        self.assertNotIn('synthetic-secret', repr(result))
        self.assertNotIn('authorName', repr(result))
        return result

    def test_mapping_and_redaction(self):
        with patch.dict(os.environ, {}, clear=True):
            self.assertEqual(self.check({'code': '0', 'data': {'verifyResult': {
                'Biren_166m': {'result': '已验证', 'records': [{'authorName': 'synthetic-secret'}]},
                'Iluvatar_bi-150': {'result': '验证中'},
                'Iluvatar_bi-100': {'result': '验证失败'},
                'unknown': {'result': '已验证'},
            }}}), ('ok', {'Biren_166m': 60027, 'Iluvatar_bi-150': 60028}))

    def test_bad_responses(self):
        with patch.dict(os.environ, {}, clear=True):
            for body, expected, raw in ((b'{bad', 'malformed_json', True),
                    ({'code': 3}, 'platform_error', False),
                    ({'code': 0, 'data': []}, 'malformed_data', False),
                    ([], 'malformed_data', False),
                    ({'code': 0, 'data': {'verifyResult': []}}, 'malformed_data', False),
                    ({'code': 0, 'data': {'verifyResult': {'Biren_166m': None}}}, 'malformed_data', False)):
                with self.subTest(body=body):
                    self.assertEqual(self.check(body, raw), (expected, {}))

    def test_transport_http_and_token_errors(self):
        for error, status in ((OSError('synthetic-secret'), 'transport_error'),
                (urllib.error.HTTPError('url', 403, 'secret', {}, None), 'http_error')):
            self.assertEqual(precheck_model('o/m', 'synthetic-secret', Mock(side_effect=error)), (status, {}))
        with patch.dict(os.environ, {}, clear=True):
            opener = Mock()
            self.assertEqual(precheck_model('o/m', opener=opener), ('missing/invalid_token', {}))
            opener.assert_not_called()
        response = Mock(status=500)
        response.__enter__ = Mock(return_value=response)
        response.__exit__ = Mock(return_value=False)
        self.assertEqual(precheck_model('o/m', 'synthetic-secret', Mock(return_value=response)), ('http_error', {}))


class TestLiveMarks(unittest.TestCase):
    def setUp(self):
        _submitted.clear()
        _marked.clear()
        self.env = patch.dict(os.environ, dict(STRATEGY_ID='sid', XC_TOKEN='synthetic-secret', CONFIG_PARAMS_SOURCE='sample'), clear=True)
        self.env.start()
        self.addCleanup(self.env.stop)
        self.addCleanup(_submitted.clear)
        self.addCleanup(_marked.clear)
        self.client = Mock()
        self.client.submit.return_value = dict(http_status=200, platform_code=0)
        self.precheck = Mock(return_value=('ok', {}))
        self.reader = Mock(return_value=(0, 'ok', set()))
        self.logs = io.StringIO()
        logpatch = patch('sys.stderr', self.logs)
        logpatch.start()
        self.addCleanup(logpatch.stop)
        # All external collaborators must be injected in this suite.
        netpatch = patch('urllib.request.urlopen', side_effect=AssertionError('network forbidden'))
        netpatch.start()
        self.addCleanup(netpatch.stop)

    def cycle(self, model='unit/model', gpus=None, rows=None):
        return run_cycle(rows or [dict(model_address=model, try_next_gpus=gpus or ['壁砺'])],
                         self.client, self.reader, precheck=self.precheck)

    def test_marks_read_failure(self):
        with patch('src.live_cycle.load_marks', side_effect=ValueError()):
            self.assertEqual(self.cycle(), [dict(action='hold', reason='marks_read_failed')])
        self.client.submit.assert_not_called()
        self.precheck.assert_not_called()

    def test_baked_mark_uses_next_choice(self):
        result = self.cycle('MadeAgents/Hammer2.1-0.5b', ['壁砺', '天垓100'])
        self.assertEqual(result[0], dict(action='skip', reason='marked:60027', model='MadeAgents/Hammer2.1-0.5b', gpu='Biren_166m'))
        self.assertEqual(self.client.submit.call_args.args[0]['targetGpu'], 'Iluvatar_bi-100')
        self.assertEqual(self.logs.getvalue().count('reason=marked:60027'), 1)

    def test_baked_mark_only_never_calls_api(self):
        result = self.cycle('MadeAgents/Hammer2.1-0.5b')
        self.assertEqual(result[-1]['reason'], 'all_gpus_done_or_unmapped')
        self.client.submit.assert_not_called()
        self.precheck.assert_not_called()
        self.reader.assert_not_called()

    def test_precheck_marks_all_remaining_choices(self):
        self.precheck.return_value = ('ok', {'Biren_166m': 60027, 'Iluvatar_bi-100': 60028})
        result = self.cycle(gpus=['壁砺', '天垓100'])
        self.assertEqual([d['reason'] for d in result], ['marked:precheck_60027', 'marked:precheck_60028', 'all_gpus_done_or_unmapped'])
        self.assertEqual(_marked, {('unit/model', 'Biren_166m'): 60027, ('unit/model', 'Iluvatar_bi-100'): 60028})
        self.cycle(gpus=['壁砺', '天垓100'])
        self.precheck.assert_called_once()
        self.client.submit.assert_not_called()

    def test_missing_mapping_submits(self):
        self.cycle()
        self.client.submit.assert_called_once()

    def test_nonblocking_platform_labels_and_missing_gpu(self):
        for label in ('验证失败', '未验证', None):
            _submitted.clear()
            response = Mock(status=200)
            mapping = {} if label is None else {'Biren_166m': {'result': label}}
            response.read.return_value = json.dumps(dict(code=0, data=dict(verifyResult=mapping))).encode()
            response.__enter__ = Mock(return_value=response)
            response.__exit__ = Mock(return_value=False)
            self.precheck.side_effect = lambda model: precheck_model(model, 'synthetic-secret', Mock(return_value=response))
            self.cycle()
        self.assertEqual(self.client.submit.call_count, 3)

    def test_precheck_uses_next_choice(self):
        self.precheck.return_value = 'ok', {'Biren_166m': 60027}
        result = self.cycle(gpus=['壁砺', '天垓100'])
        self.assertEqual(result[0]['reason'], 'marked:precheck_60027')
        self.assertEqual(self.client.submit.call_args.args[0]['targetGpu'], 'Iluvatar_bi-100')
        self.precheck.assert_called_once()

    def test_precheck_failure_and_unsafe_status(self):
        for status in ('transport_error', 'synthetic-secret authorName'):
            with self.subTest(status=status):
                _submitted.clear()
                self.precheck.return_value = status, {}
                result = self.cycle()
                self.assertEqual(result[0]['action'], 'precheck')
                self.assertEqual(result[-1]['action'], 'submitted')
        self.assertNotIn('synthetic-secret', self.logs.getvalue())
        self.assertNotIn('authorName', self.logs.getvalue())

    def test_disabled(self):
        with patch.dict(os.environ, {'LIVE_PRECHECK': '0'}):
            self.cycle()
        self.precheck.assert_not_called()
        self.client.submit.assert_called_once()

    def test_precheck_cap(self):
        self.precheck.return_value = 'ok', {'Biren_166m': 60027}
        rows = [dict(model_address=f'unit/m{i}', targetGpu='壁砺') for i in range(32)]
        for value, count in (('2', 2), ('bad', 30), ('0', 30)):
            _marked.clear()
            self.precheck.reset_mock()
            with patch.dict(os.environ, {'LIVE_MAX_PRECHECKS_PER_CYCLE': value, 'LIVE_MAX_ATTEMPTS_PER_CYCLE': '1'}):
                result = self.cycle(rows=rows)
            self.assertEqual(self.precheck.call_count, count)
            self.assertEqual(result[-1]['reason'], 'precheck_cap')
            self.assertEqual(sum(d['reason'] == 'precheck_cap' for d in result), 1)
        self.client.submit.assert_not_called()

    def test_refresh_reuses_precheck_mapping(self):
        self.precheck.return_value = 'ok', {'Iluvatar_bi-100': 60028}
        self.reader.return_value = 0, 'ok', {('unit/model', 'Biren_166m')}
        result = self.cycle(gpus=['壁砺', '天垓100'])
        self.precheck.assert_called_once()
        self.client.submit.assert_not_called()
        self.assertEqual(result[-1]['reason'], 'all_gpus_done_or_unmapped')
        self.assertEqual(sum(d['reason'] == 'marked:precheck_60028' for d in result), 1)

    def test_rejection_memory_and_exact_log(self):
        self.client.submit.return_value = dict(http_status=200, platform_code=60022)
        self.cycle()
        result = self.cycle()
        self.client.submit.assert_called_once()
        self.assertEqual(result[0]['reason'], 'marked:60022')
        self.assertIn('[live] rejected model=unit/model gpu=Biren_166m code=60022 http=200 queue_before=0 reason=CREATE_CONTEST_TASK_GPU_CONFIG_UNSUPPORTED\n', self.logs.getvalue())

    def test_no_state_file_writes(self):
        with tempfile.TemporaryDirectory() as directory:
            path = Path(directory) / 'state.json'
            with patch.dict(os.environ, {'LIVE_STATE_FILE': str(path)}), patch.object(Path, 'write_text', side_effect=AssertionError('file write forbidden')):
                self.cycle()
            self.assertFalse(path.exists())


class TestSyncMarks(unittest.TestCase):
    def test_sync_and_dry_run(self):
        with tempfile.TemporaryDirectory() as directory:
            marks = Path(directory) / 'marks.json'
            log = Path(directory) / 'log.txt'
            existing = dict(model='Org/Existing', gpu='Biren_166m', code=60027, reason='first reason', date='2020-01-01')
            marks.write_text(json.dumps(dict(version=1, marks=[existing])))
            log.write_text('''timestamp console [live] rejected model=org/existing gpu=Biren_166m code=60028 http=200
prefix [live] rejected model=Org/New gpu=Iluvatar_bi-150 code=60022 http=200
[live] rejected model=org/new gpu=Iluvatar_bi-150 code=60027
[live] skip model=Org/Precheck gpu=Biren_166m reason=marked:precheck_60028
[live] rejected model=Org/Limit gpu=Biren_166m code=60007
[live] rejected model=Org/Disabled gpu=Biren_166m code=60014
[live] rejected model=Org/Success gpu=Biren_166m code=0
[live] rejected model=invalid gpu=Biren_166m code=60022
[live] rejected model=Org/Invalid gpu=unknown code=60022
''')
            command = [sys.executable, 'scripts/sync_reject_marks.py', str(log), '--marks', str(marks), '--date', '2026-10-17']
            before = marks.read_bytes()
            result = subprocess.run(command + ['--dry-run'], capture_output=True, text=True, check=True)
            self.assertEqual(result.stdout, 'added 2 total 3\n')
            self.assertEqual(result.stderr, '')
            self.assertEqual(marks.read_bytes(), before)
            result = subprocess.run(command, capture_output=True, text=True, check=True)
            self.assertEqual(result.stdout, 'added 2 total 3\n')
            rows = json.loads(marks.read_text())['marks']
            self.assertEqual(rows[0], existing)
            self.assertEqual(rows[1]['model'], 'Org/New')
            self.assertEqual(rows[1]['code'], 60022)
            self.assertEqual(rows[2]['code'], 60028)
            self.assertTrue(marks.read_bytes().endswith(b'\n'))
            self.assertEqual(len(load_marks(marks)), 3)
            result = subprocess.run(command, capture_output=True, text=True, check=True)
            self.assertEqual(result.stdout, 'added 0 total 3\n')