import sys import unittest from pathlib import Path from unittest.mock import patch sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "modelhub_submmit_api")) from modelhub_client import ModelHubClient, ModelHubClientPool, ModelHubAPIError class CapacityRefreshTests(unittest.TestCase): def test_api_boundary_clamps_all_callers(self): client = ModelHubClient(token="test") with patch.object(client, "_request", return_value={}) as request: client.list_tasks_page(page_size=200) self.assertEqual(100, request.call_args.kwargs["query"]["pageSize"]) def client(self, total=97, observed=94): client = ModelHubClient(token="test") def page(*, status, page_size, **kwargs): self.assertLessEqual(page_size, 100) return {"data": {"total": total if status == "waiting" else 0, "records": [ {"taskId": str(i), "status": status} for i in range(observed if status == "waiting" else 0)]}} client.list_tasks_page = page return client def test_full_cache_refreshes_to_free_slots_with_disagreeing_total(self): client = self.client() pool = ModelHubClientPool([client], active_task_cap=100) pool._remote_counts[0] = 100 self.assertEqual([97], pool.refresh_active_counts()) self.assertEqual(3, pool.available_submit_slots()) self.assertTrue(client.active_count_details[0]["mismatch"]) def test_error_is_unknown_and_recovery_reopens_slots(self): client = self.client() pool = ModelHubClientPool([client], active_task_cap=100) pool.refresh_active_counts() with patch.object(client, "list_tasks_page", side_effect=ModelHubAPIError("分页大小不能超过100")): pool.refresh_active_counts() self.assertEqual(0, pool.available_submit_slots()) self.assertIsNone(pool.active_count_status()[0]["active"]) self.assertIsNone(pool._reserve_account(set())) pool.refresh_active_counts() self.assertEqual(3, pool.available_submit_slots()) def test_incomplete_or_repeated_pages_do_not_open_slots(self): client = self.client(total=200, observed=100) with self.assertRaises(ModelHubAPIError): client.count_active_tasks_exact() client.list_tasks_page = lambda **kwargs: {"data": {"records": []}} with self.assertRaises(ModelHubAPIError): client.count_active_tasks_exact() def test_duplicate_ids_and_transition_are_counted_once(self): client = ModelHubClient(token="test") client.list_tasks_page = lambda **kwargs: {"data": {"total": 1, "records": [ {"taskId": "same", "status": kwargs["status"]}, {"taskId": "same", "status": kwargs["status"]}]}} self.assertEqual(1, client.count_active_tasks_exact()) def test_total_lower_than_observed_uses_larger_count(self): self.assertEqual(99, self.client(total=97, observed=99).count_active_tasks_exact())