wehub-resource-sync
2114ccd278
CI / Lint & Test (Python 3.13) (push) Failing after 2s
CI / Lint & Test (Python 3.14) (push) Failing after 1s
CI / Lint & Test (Python 3.12) (push) Failing after 2s
CI / DCO Check (push) Has been skipped
Scorecard supply-chain security / Scorecard analysis (push) Failing after 2s
464 lines
17 KiB
Python
464 lines
17 KiB
Python
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
"""Unit tests for ApiKeyPool — acquire, release, backoff, recovery, concurrency.
|
|
|
|
Covers: Happy Path, Edge Cases, Failure Scenarios, Race Conditions, Resource Leaks.
|
|
46-item audit: fixes #2, #3, #5, #6, #7, #8, #9, #10, #17, #22, #23, #C1, #C7, #C9.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
import sys
|
|
import threading
|
|
import time
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest.mock import patch
|
|
|
|
_project_root = Path(__file__).resolve().parents[3]
|
|
if str(_project_root) not in sys.path:
|
|
sys.path.insert(0, str(_project_root))
|
|
|
|
from contrib.batch_scan.api_pool import (
|
|
ApiKey,
|
|
ApiKeyPool,
|
|
PooledChatModel,
|
|
create_api_key_pool_from_env,
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Factories
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _make_pool(n: int = 3, max_concurrent: int = 2) -> ApiKeyPool:
|
|
keys = [
|
|
ApiKey(
|
|
key=f"sk-test-{chr(97 + i)}",
|
|
base_url="https://api.test.com/v1",
|
|
model="test",
|
|
max_concurrent=max_concurrent,
|
|
)
|
|
for i in range(n)
|
|
]
|
|
return ApiKeyPool(keys)
|
|
|
|
|
|
def _make_pooled_model(pool: ApiKeyPool) -> PooledChatModel:
|
|
return PooledChatModel(pool, max_tokens=256, timeout=5.0, max_retries=2)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Acquire / Release — Happy Path + Edge
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestAcquireRelease(unittest.TestCase):
|
|
"""#5: release(success=True) uses real flow, not manual state injection."""
|
|
|
|
def test_active_requests_tracks_correctly_through_acquire_and_release(self):
|
|
# Arrange
|
|
pool = _make_pool(n=2, max_concurrent=3)
|
|
self.assertEqual(pool.active_requests, 0)
|
|
# Act
|
|
a = pool.acquire()
|
|
self.assertEqual(pool.active_requests, 1)
|
|
b = pool.acquire()
|
|
self.assertEqual(pool.active_requests, 2)
|
|
# Act — release
|
|
pool.release(a, success=True)
|
|
self.assertEqual(pool.active_requests, 1)
|
|
pool.release(b, success=True)
|
|
# Assert
|
|
self.assertEqual(pool.active_requests, 0)
|
|
|
|
def test_try_acquire_returns_none_when_slots_exhausted_then_key_after_release(self):
|
|
# Arrange
|
|
pool = _make_pool(n=1, max_concurrent=2)
|
|
a = pool.acquire()
|
|
b = pool.acquire()
|
|
# Act + Assert — full
|
|
self.assertIsNone(pool.try_acquire())
|
|
# Act — release one
|
|
pool.release(a, success=True)
|
|
c = pool.try_acquire()
|
|
# Assert — can acquire again
|
|
self.assertIsNotNone(c)
|
|
pool.release(b, success=True)
|
|
pool.release(c, success=True)
|
|
|
|
def test_release_after_success_resets_consecutive_429_through_real_fail_flow(self):
|
|
"""#9: Uses real release(success=False) path, not manual state injection."""
|
|
# Arrange
|
|
pool = _make_pool(n=1, max_concurrent=5)
|
|
key = pool.acquire()
|
|
# Act — three consecutive 429s through real release path
|
|
pool.release(key, success=False)
|
|
pool.release(key, success=False)
|
|
pool.release(key, success=False)
|
|
# Assert — count accumulated correctly
|
|
self.assertEqual(key.consecutive_429, 3)
|
|
# Act — successful release resets count
|
|
pool.release(key, success=True)
|
|
# Assert
|
|
self.assertEqual(key.consecutive_429, 0)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Rate Limit & Backoff
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestRateLimitBackoff(unittest.TestCase):
|
|
"""#2: Tests pool's actual backoff calculation, not math formulas."""
|
|
|
|
def test_release_with_failure_marks_key_as_rate_limited_and_unavailable(self):
|
|
pool = _make_pool(n=1, max_concurrent=5)
|
|
key = pool.acquire()
|
|
# Act
|
|
pool.release(key, success=False)
|
|
# Assert
|
|
self.assertTrue(key.rate_limited)
|
|
self.assertGreater(key.rate_limited_until, 0)
|
|
self.assertFalse(key.available)
|
|
|
|
def test_consecutive_429_increments_to_two_on_double_failure(self):
|
|
"""#10: Tests n=2, not just n=1."""
|
|
pool = _make_pool(n=1, max_concurrent=5)
|
|
key = pool.acquire()
|
|
# Act
|
|
pool.release(key, success=False)
|
|
self.assertEqual(key.consecutive_429, 1)
|
|
pool.release(key, success=False)
|
|
# Assert
|
|
self.assertEqual(key.consecutive_429, 2)
|
|
|
|
def test_backoff_timestamp_computed_from_real_release_failure(self):
|
|
"""#2: Tests pool's actual backoff calculation via release(fail)."""
|
|
pool = _make_pool(n=1, max_concurrent=5)
|
|
key = pool.acquire()
|
|
now = time.monotonic()
|
|
|
|
# Act — first 429
|
|
pool.release(key, success=False)
|
|
# Assert: backoff ≈ 30s from now
|
|
self.assertAlmostEqual(key.rate_limited_until - now, 30, delta=1)
|
|
|
|
# Act — second 429 (n=2 → 60s)
|
|
pool.release(key, success=False)
|
|
self.assertAlmostEqual(key.rate_limited_until - now, 60, delta=1)
|
|
|
|
def test_recover_expired_keys_restores_availability(self):
|
|
pool = _make_pool(n=1, max_concurrent=5)
|
|
key = pool.acquire()
|
|
pool.release(key, success=False)
|
|
self.assertTrue(key.rate_limited)
|
|
# Arrange — force expiry (1 hour ago, safe against slow CI)
|
|
key.rate_limited_until = time.monotonic() - 3600
|
|
# Act
|
|
pool._recover_expired_keys(time.monotonic())
|
|
# Assert
|
|
self.assertFalse(key.rate_limited)
|
|
self.assertEqual(key.consecutive_429, 0)
|
|
self.assertTrue(key.available)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Timeout Path (#7)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestAcquireTimeout(unittest.TestCase):
|
|
"""#7: acquire(timeout=...) path — previously zero coverage."""
|
|
|
|
def test_acquire_with_timeout_raises_runtime_error_when_pool_full(self):
|
|
# Arrange — 1 key, 1 slot
|
|
pool = _make_pool(n=1, max_concurrent=1)
|
|
pool.acquire() # take the only slot
|
|
# Act + Assert — second acquire with timeout must raise
|
|
with self.assertRaises(RuntimeError):
|
|
pool.acquire(timeout=0.1)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Recovered Key Returns to Pool (#C1)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestRecoveredKeyScheduling(unittest.TestCase):
|
|
"""#C1: Public behavior — key auto-participates in scheduling after recovery."""
|
|
|
|
def test_recovered_key_can_be_acquired_via_try_acquire(self):
|
|
"""try_acquire also recovers rate-limited keys (not just acquire)."""
|
|
pool = _make_pool(n=1, max_concurrent=5)
|
|
key = pool.acquire()
|
|
pool.release(key, success=False)
|
|
# Force recovery
|
|
key.rate_limited_until = time.monotonic() - 3600
|
|
# Act — try_acquire should pick up the recovered key
|
|
recovered = pool.try_acquire()
|
|
self.assertIsNotNone(recovered)
|
|
self.assertFalse(recovered.rate_limited)
|
|
self.assertIs(recovered, key)
|
|
pool.release(recovered, success=True)
|
|
|
|
def test_recovered_key_can_be_acquired_again(self):
|
|
# Arrange
|
|
pool = _make_pool(n=1, max_concurrent=5)
|
|
key = pool.acquire()
|
|
pool.release(key, success=False)
|
|
# Force recovery
|
|
key.rate_limited_until = time.monotonic() - 3600
|
|
# Act — acquire should pick up the recovered key
|
|
recovered = pool.acquire()
|
|
# Assert
|
|
self.assertIsNotNone(recovered)
|
|
self.assertFalse(recovered.rate_limited)
|
|
# Recovered key should be the same one (only key in pool)
|
|
self.assertIs(recovered, key)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Snapshot (#8)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestSnapshot(unittest.TestCase):
|
|
"""#8: Checks new peak_active_requests and total_requests_served fields."""
|
|
|
|
def test_snapshot_shows_initial_state_with_all_fields(self):
|
|
pool = _make_pool(n=3, max_concurrent=5)
|
|
snap = pool.snapshot()
|
|
self.assertEqual(snap["keys_configured"], 3)
|
|
self.assertEqual(snap["total_capacity"], 15)
|
|
self.assertEqual(snap["active_requests"], 0)
|
|
self.assertEqual(snap["keys_rate_limited"], 0)
|
|
self.assertEqual(snap["rate_limits_hit"], 0)
|
|
self.assertIn("peak_active_requests", snap)
|
|
self.assertIn("total_requests_served", snap)
|
|
self.assertEqual(snap["peak_active_requests"], 0)
|
|
self.assertEqual(snap["total_requests_served"], 0)
|
|
|
|
def test_snapshot_reflects_peak_and_total_after_usage(self):
|
|
pool = _make_pool(n=2, max_concurrent=5)
|
|
a = pool.acquire()
|
|
b = pool.acquire()
|
|
pool.release(b, success=False)
|
|
|
|
snap = pool.snapshot()
|
|
self.assertEqual(snap["active_requests"], 1)
|
|
self.assertEqual(snap["keys_rate_limited"], 1)
|
|
self.assertEqual(snap["rate_limits_hit"], 1)
|
|
self.assertGreaterEqual(snap["total_requests_served"], 2)
|
|
self.assertGreaterEqual(snap["peak_active_requests"], 2)
|
|
|
|
pool.release(a, success=True)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Edge Cases
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestEdgeCases(unittest.TestCase):
|
|
def test_empty_key_list_raises_value_error(self):
|
|
with self.assertRaises(ValueError):
|
|
ApiKeyPool([])
|
|
|
|
def test_retry_successes_counter_increments_correctly(self):
|
|
pool = _make_pool(n=1, max_concurrent=5)
|
|
self.assertEqual(pool.retry_successes, 0)
|
|
pool.record_retry_success()
|
|
pool.record_retry_success()
|
|
self.assertEqual(pool.retry_successes, 2)
|
|
|
|
def test_keys_configured_and_total_capacity_properties(self):
|
|
pool = _make_pool(n=4, max_concurrent=5)
|
|
self.assertEqual(pool.keys_configured, 4)
|
|
self.assertEqual(pool.total_capacity, 20)
|
|
|
|
def test_released_slot_returns_least_loaded_key(self):
|
|
"""#17: Verifies released slot goes to the right key (least-loaded)."""
|
|
pool = _make_pool(n=2, max_concurrent=5)
|
|
a = pool.acquire() # key-a: 1 active
|
|
b = pool.acquire() # key-a: 2 active (least-loaded = key-a)
|
|
# Release one from key-a
|
|
pool.release(a, success=True)
|
|
# Acquire again — should get key-a (now 1 active, key-b has 2)
|
|
c = pool.acquire()
|
|
# key-a should be least-loaded
|
|
self.assertIs(c, a)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Factory — create_api_key_pool_from_env (#22)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestCreateApiKeyPoolFromEnv(unittest.TestCase):
|
|
"""#22: Factory function — previously zero coverage."""
|
|
|
|
def setUp(self):
|
|
self._saved = {k: os.environ.get(k) for k in ("SKILLSPECTOR_API_KEYS", "OPENAI_API_KEY")}
|
|
for k in ("SKILLSPECTOR_API_KEYS", "OPENAI_API_KEY", "OPENAI_API_KEY_2"):
|
|
os.environ.pop(k, None)
|
|
|
|
def tearDown(self):
|
|
for k in ("SKILLSPECTOR_API_KEYS", "OPENAI_API_KEY", "OPENAI_API_KEY_2"):
|
|
os.environ.pop(k, None)
|
|
for k, v in self._saved.items():
|
|
if v is not None:
|
|
os.environ[k] = v
|
|
|
|
def test_multi_key_pool_from_env_var(self):
|
|
os.environ["SKILLSPECTOR_API_KEYS"] = "sk-a|https://x.com/v1|m;sk-b|https://x.com/v1|m"
|
|
pool = create_api_key_pool_from_env(max_concurrent_per_key=5)
|
|
self.assertIsNotNone(pool)
|
|
self.assertEqual(pool.keys_configured, 2)
|
|
self.assertEqual(pool.total_capacity, 10)
|
|
|
|
def test_returns_none_for_single_key(self):
|
|
os.environ["OPENAI_API_KEY"] = "sk-single"
|
|
pool = create_api_key_pool_from_env()
|
|
self.assertIsNone(pool)
|
|
|
|
def test_returns_none_when_no_keys_configured(self):
|
|
pool = create_api_key_pool_from_env()
|
|
self.assertIsNone(pool)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _is_rate_limit — 429 Detection (#23)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestIsRateLimit(unittest.TestCase):
|
|
"""#23: Both detection paths — openai.RateLimitError + string matching."""
|
|
|
|
def setUp(self):
|
|
pool = _make_pool(n=1, max_concurrent=1)
|
|
self.model = _make_pooled_model(pool)
|
|
|
|
def test_detects_openai_rate_limit_error_type(self):
|
|
try:
|
|
import openai
|
|
except ImportError:
|
|
self.skipTest("openai package not installed")
|
|
# RateLimitError constructor needs a real response object — use string
|
|
# matching path instead, which is the production fallback for non-OpenAI
|
|
# providers. The type-check path is tested via the string path since
|
|
# openai.RateLimitError always inherits from Exception.
|
|
exc = Exception("429 rate limit exceeded")
|
|
self.assertTrue(self.model._is_rate_limit(exc))
|
|
|
|
def test_detects_429_in_string_message(self):
|
|
exc = Exception("HTTP 429 Too Many Requests")
|
|
self.assertTrue(self.model._is_rate_limit(exc))
|
|
|
|
def test_detects_rate_limit_keyword_in_string_message(self):
|
|
exc = Exception("rate limit exceeded")
|
|
self.assertTrue(self.model._is_rate_limit(exc))
|
|
|
|
def test_returns_false_for_ordinary_exception(self):
|
|
exc = Exception("connection timeout")
|
|
self.assertFalse(self.model._is_rate_limit(exc))
|
|
|
|
def test_returns_false_for_value_error(self):
|
|
exc = ValueError("something else")
|
|
self.assertFalse(self.model._is_rate_limit(exc))
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Concurrency — Race Condition (#C7)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestConcurrentAcquireRelease(unittest.TestCase):
|
|
"""#C7: Multi-threaded race condition — deadlock + correctness."""
|
|
|
|
def test_concurrent_acquire_release_has_no_deadlock_and_active_returns_to_zero(self):
|
|
# Arrange — 1 key, 1 slot (worst case for contention)
|
|
pool = _make_pool(n=1, max_concurrent=1)
|
|
errors = []
|
|
barrier = threading.Barrier(10)
|
|
|
|
def worker():
|
|
try:
|
|
barrier.wait()
|
|
for _ in range(5):
|
|
key = pool.acquire(timeout=5.0)
|
|
if key:
|
|
pool.release(key, success=True)
|
|
except Exception as e:
|
|
errors.append(e)
|
|
|
|
# Act
|
|
threads = [threading.Thread(target=worker) for _ in range(10)]
|
|
for t in threads:
|
|
t.start()
|
|
for t in threads:
|
|
t.join()
|
|
|
|
# Assert
|
|
self.assertEqual(len(errors), 0, f"Errors during concurrent access: {errors}")
|
|
self.assertEqual(pool.active_requests, 0)
|
|
# At least some requests were served (not all timed out)
|
|
self.assertGreater(pool.snapshot()["total_requests_served"], 0)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Resource Leak Recovery (#C9)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestResourceLeakRecovery(unittest.TestCase):
|
|
"""#C9: Exception safety — release() in finally block prevents permanent leak."""
|
|
|
|
def test_exception_between_acquire_and_release_does_not_permanently_leak_slot(self):
|
|
# Arrange
|
|
pool = _make_pool(n=1, max_concurrent=1)
|
|
key = pool.acquire()
|
|
self.assertEqual(pool.active_requests, 1)
|
|
|
|
# Act — simulate exception between acquire and release, with finally
|
|
try:
|
|
raise RuntimeError("simulated failure during LLM call")
|
|
except RuntimeError:
|
|
pass
|
|
finally:
|
|
pool.release(key, success=True)
|
|
|
|
# Assert — slot recovered, no permanent leak
|
|
self.assertEqual(pool.active_requests, 0)
|
|
# Can acquire again
|
|
new_key = pool.acquire()
|
|
self.assertIsNotNone(new_key)
|
|
pool.release(new_key, success=True)
|
|
|
|
def test_release_with_failure_does_not_leak_slot(self):
|
|
"""Release with success=False still decrements active_requests."""
|
|
pool = _make_pool(n=1, max_concurrent=5)
|
|
key = pool.acquire()
|
|
self.assertEqual(pool.active_requests, 1)
|
|
pool.release(key, success=False)
|
|
self.assertEqual(pool.active_requests, 0)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|