submmit/tests/test_capacity_refresh.py

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, current=1, **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(
                    (current - 1) * page_size, min(current * page_size, 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([94], pool.refresh_active_counts())
        self.assertEqual(6, 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(6, pool.available_submit_slots())

    def test_incomplete_or_repeated_pages_do_not_open_slots(self):
        client = self.client(total=200, observed=100)
        client.list_tasks_page = lambda **kwargs: {"data": {"records": [
            {"taskId": "same", "status": kwargs["status"]}]}}
        with self.assertRaises(ModelHubAPIError):
            client.count_active_tasks_exact()
        client.list_tasks_page = lambda **kwargs: {"data": {}}
        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": [] if kwargs['current'] > 1 else [
            {"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())

    def test_inflated_total_does_not_hide_capacity(self):
        self.assertEqual(100, self.client(total=773, observed=100).count_active_tasks_exact())

    def test_underreported_pages_do_not_truncate_scan(self):
        client = self.client(total=58, observed=175)
        original = client.list_tasks_page
        def page(**kwargs):
            result = original(**kwargs)
            result['data']['pages'] = 1
            return result
        client.list_tasks_page = page
        self.assertEqual(175, client.count_active_tasks_exact())

    def test_different_confirmation_ids_fail_closed_even_with_equal_counts(self):
        client = self.client(total=1, observed=1)
        original = client.list_tasks_page
        def page(**kwargs):
            result = original(**kwargs)
            if kwargs['page_size'] == 50:
                for row in result['data']['records']:
                    row['taskId'] = 'replacement'
            return result
        client.list_tasks_page = page
        with self.assertRaisesRegex(ModelHubAPIError, 'active_snapshot_changed'):
            client.count_active_tasks_exact()

    def test_confirmation_error_does_not_open_slots(self):
        client = self.client(total=773, observed=94)
        original = client.list_tasks_page
        def page(**kwargs):
            if kwargs['page_size'] == 50:
                raise ModelHubAPIError('temporary failure')
            return original(**kwargs)
        client.list_tasks_page = page
        pool = ModelHubClientPool([client], active_task_cap=100)
        pool.refresh_active_counts()
        self.assertEqual(0, pool.available_submit_slots())
        self.assertFalse(pool.active_count_status()[0]['known'])

    def test_empty_queue_confirmed_despite_stale_total(self):
        self.assertEqual(0, self.client(total=773, observed=0).count_active_tasks_exact())

    def test_short_page_is_not_end_of_scan(self):
        client = ModelHubClient(token='test')
        def page(**kwargs):
            rows = [{'taskId': str(kwargs['current']), 'status': 'waiting'}] if (
                kwargs['status'] == 'waiting' and kwargs['current'] <= 3) else []
            return {'data': {'records': rows}}
        client.list_tasks_page = page
        self.assertEqual(3, client.count_active_tasks_exact())

    def test_scan_budget_is_bounded(self):
        client = ModelHubClient(token='test')
        client.list_tasks_page = lambda **kw: {'data': {'records': [
            {'taskId': str(kw['current']), 'status': kw['status']}]}}
        with self.assertRaisesRegex(ModelHubAPIError, 'active_pagination_budget_exhausted'):
            client.count_active_tasks_exact()