import io import os import tempfile import unittest from unittest.mock import Mock, patch from src.live_cycle import run_cycle, _submitted def candidate(name, section='B', foothold='unknown'): return dict(model_address='org/' + name, section=section, foothold_class=foothold, try_next_gpus=['壁砺']) class TestLiveCycle(unittest.TestCase): def setUp(self): directory = tempfile.TemporaryDirectory() self.addCleanup(directory.cleanup) self.env = patch.dict(os.environ, {'STRATEGY_ID': 'sid', 'LIVE_STATE_FILE': directory.name + '/state.json', 'CONFIG_PARAMS_SOURCE': 'sample', 'XC_TOKEN': 'synthetic-secret'}, clear=True) self.env.start() self.addCleanup(self.env.stop) _submitted.clear() self.client = Mock() self.client.submit.return_value = dict(http_status=200, platform_code=0) network = patch('urllib.request.urlopen', side_effect=AssertionError('network forbidden')) network.start() self.addCleanup(network.stop) logs = patch('sys.stderr', new_callable=io.StringIO) self.logs = logs.start() self.addCleanup(logs.stop) def cycle(self, rows=None, counts=(94,)): reader = Mock(side_effect=[(n, 'platform_error' if n is None else 'ok', set()) for n in counts]) return run_cycle(rows or [candidate('a')], self.client, reader) def test_high_and_failed_queue(self): for n in (95, None): decisions = self.cycle(counts=(n,)) self.assertEqual(decisions[-1]['action'], 'hold') self.client.submit.assert_not_called() self.assertIn('[live] hold', self.logs.getvalue()) def test_order_and_cap(self): rows = [candidate('b'), candidate('a', 'A'), candidate('verified', 'A', 'verified_reuse'), candidate('c'), candidate('d')] self.cycle(rows, (90, 90, 90)) self.assertEqual([c.args[0]['modelAddress'] for c in self.client.submit.call_args_list], ['org/verified', 'org/a', 'org/b']) self.assertIn('cycle_cap', self.logs.getvalue()) def test_midcycle(self): self.cycle([candidate(str(i)) for i in range(5)], (93, 94, 95)) self.assertEqual(self.client.submit.call_count, 2) def test_kill_and_threshold(self): for env in ({'ALLOW_LIVE_SUBMIT': '0'}, {'LIVE_QUEUE_THRESHOLD': 'bad'}, {'LIVE_QUEUE_THRESHOLD': '20'}): with patch.dict(os.environ, env): self.cycle(counts=(20,)) self.client.submit.assert_not_called() with patch.dict(os.environ, {'LIVE_QUEUE_THRESHOLD': '100'}): self.cycle(counts=(99,)) self.client.submit.assert_called_once() def test_banned_duplicate_persistence(self): rows = [candidate('GPT-2'), candidate('ok'), candidate('ok')] self.cycle(rows) _submitted.clear() self.cycle([candidate('ok')], ()) self.client.submit.assert_called_once() self.assertIn('banned:', self.logs.getvalue()) self.assertIn('all_gpus_done_or_unmapped', self.logs.getvalue()) def test_remote_duplicate_and_transport_reservation(self): run_cycle([candidate('remote')], self.client, lambda: (1, 'ok', {('org/remote', 'Biren_166m')})) self.client.submit.assert_not_called() self.client.submit.side_effect = OSError('do not log') self.cycle() self.cycle(counts=()) self.client.submit.assert_called_once() self.assertNotIn('do not log', self.logs.getvalue()) def test_missing_token_real_reader(self): with patch.dict(os.environ, {'XC_TOKEN': ''}): run_cycle([candidate('a')], self.client) self.client.submit.assert_not_called() self.assertIn('missing/invalid_token', self.logs.getvalue()) def test_worker_waits_and_stops(self): from src.server import _live_worker with patch('src.server._shutdown') as shutdown, patch('src.server.run_cycle') as cycle: shutdown.is_set.return_value = False shutdown.wait.return_value = True with patch.dict(os.environ, {'LIVE_CYCLE_SECONDS': '1'}): _live_worker() cycle.assert_called_once() shutdown.wait.assert_called_once_with(60) def test_unauthenticated_platform_zero_submits(self): import json response = Mock() response.__enter__ = Mock(return_value=response) response.__exit__ = Mock(return_value=False) response.status = 200 response.read.return_value = json.dumps({'code': 40100}).encode() with patch.dict(os.environ, {'XC_TOKEN': 'synthetic-secret'}), patch('urllib.request.urlopen', return_value=response): run_cycle([candidate('a')], self.client) self.client.submit.assert_not_called() def test_unmapped_and_placeholder(self): rows = [dict(model_address='org/unmapped', targetGpu='unknown'), candidate('')] self.cycle(rows, ()) self.client.submit.assert_not_called() self.assertIn('all_gpus_done_or_unmapped', self.logs.getvalue()) self.assertIn('placeholder', self.logs.getvalue())