sashatrask
140 lines
5.2 KiB
Python
140 lines
5.2 KiB
Python
from pathlib import Path
|
|
import sys
|
|
import unittest
|
|
|
|
|
|
ROOT = Path(__file__).resolve().parents[1]
|
|
APP_DIR = ROOT / 'app'
|
|
sys.path.insert(0, str(APP_DIR))
|
|
|
|
import scanner
|
|
from keycheck_candidates import extract_candidates
|
|
from keycheckers import provider_resolution
|
|
from keycheckers.provider_resolver import providerResolverKeycheck
|
|
|
|
|
|
class AmbiguousProviderResolutionTests(unittest.TestCase):
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
scanner.initialize_scanner_runtime(preflight_complete=True, register_cleanup=False)
|
|
|
|
def test_weak_generic_key_routes_to_one_resolver_candidate(self):
|
|
key = 'sk-' + ('a' * 32)
|
|
finding = {'DetectorName': 'DeepSeek', 'Raw': key}
|
|
|
|
scanner.strip_nearby_context_for_persistence({'findings': [finding]})
|
|
candidates = list(extract_candidates(finding))
|
|
|
|
self.assertEqual(finding['ScannerContext']['provider_hint'], 'ambiguous_generic_sk')
|
|
self.assertEqual(
|
|
finding['ScannerContext']['provider_candidates'],
|
|
['deepseek', 'zai', 'qwen', 'kimi'],
|
|
)
|
|
self.assertEqual(len(candidates), 1)
|
|
self.assertEqual(candidates[0].service, 'provider_resolver')
|
|
|
|
def test_resolver_can_recover_a_misrouted_dotted_zai_key(self):
|
|
key = ('a' * 32) + '.' + ('B' * 16)
|
|
self.assertEqual(providerResolverKeycheck.key_rejection_reason(key), '')
|
|
|
|
def test_strong_context_keeps_direct_provider_route(self):
|
|
key = 'sk-' + ('b' * 32)
|
|
finding = {
|
|
'DetectorName': 'DeepSeek',
|
|
'Raw': key,
|
|
'ScannerContext': {'nearby': 'DEEPSEEK_API_KEY is configured here'},
|
|
}
|
|
|
|
scanner.strip_nearby_context_for_persistence({'findings': [finding]})
|
|
candidates = list(extract_candidates(finding))
|
|
|
|
self.assertEqual(finding['ScannerContext']['provider_hint'], 'deepseek')
|
|
self.assertNotIn('provider_candidates', finding['ScannerContext'])
|
|
self.assertEqual([candidate.service for candidate in candidates], ['deepseek'])
|
|
|
|
def test_provider_order_is_origin_first_and_bounded_by_evidence(self):
|
|
finding = {
|
|
'DetectorName': 'Qwen',
|
|
'ScannerContext': {
|
|
'provider_hint': 'ambiguous_generic_sk',
|
|
'provider_candidates': ['zai', 'qwen'],
|
|
},
|
|
}
|
|
self.assertEqual(
|
|
provider_resolution.ordered_providers(finding, origin_service='qwen'),
|
|
['qwen', 'zai'],
|
|
)
|
|
|
|
def test_first_match_short_circuits_remaining_providers(self):
|
|
calls = []
|
|
|
|
def probe(provider, _key, _proxy, _timeout, _debug):
|
|
calls.append(provider)
|
|
if provider == 'zai':
|
|
return {'status': 'VALID', 'authenticated': True, 'model_count': 9}
|
|
return {'status': 'DEAD', 'http_status': 401}
|
|
|
|
result = provider_resolution.resolve_provider_key(
|
|
'sk-' + ('c' * 32),
|
|
hint='ambiguous_generic_sk',
|
|
configured_order=['deepseek', 'zai', 'qwen', 'kimi'],
|
|
probe=probe,
|
|
)
|
|
|
|
self.assertEqual(calls, ['deepseek', 'zai'])
|
|
self.assertEqual(result['provider_resolution'], 'matched')
|
|
self.assertEqual(result['resolved_provider'], 'zai')
|
|
self.assertEqual(
|
|
[attempt['outcome'] for attempt in result['provider_resolution_attempts']],
|
|
['no_match', 'match'],
|
|
)
|
|
|
|
def test_retry_does_not_block_a_later_positive_match(self):
|
|
calls = []
|
|
|
|
def probe(provider, _key, _proxy, _timeout, _debug):
|
|
calls.append(provider)
|
|
if provider == 'deepseek':
|
|
return {'status': 'NETWORK', 'message': 'timeout'}
|
|
if provider == 'zai':
|
|
return {'status': 'NO_BALANCE', 'authenticated': True}
|
|
return {'status': 'DEAD'}
|
|
|
|
result = provider_resolution.resolve_provider_key(
|
|
'sk-' + ('d' * 32), hint='ambiguous_generic_sk', probe=probe,
|
|
)
|
|
|
|
self.assertEqual(calls, ['deepseek', 'zai'])
|
|
self.assertEqual(result['resolved_provider'], 'zai')
|
|
self.assertEqual(result['status'], 'NO_BALANCE')
|
|
|
|
def test_inconclusive_route_is_not_exhausted(self):
|
|
def probe(provider, _key, _proxy, _timeout, _debug):
|
|
if provider == 'deepseek':
|
|
return {'status': 'NETWORK', 'message': 'timeout'}
|
|
return {'status': 'DEAD'}
|
|
|
|
result = provider_resolution.resolve_provider_key(
|
|
'sk-' + ('e' * 32), hint='ambiguous_generic_sk', probe=probe,
|
|
)
|
|
|
|
self.assertEqual(result['status'], 'NETWORK')
|
|
self.assertEqual(result['provider_resolution'], 'retry')
|
|
self.assertNotIn('resolved_provider', result)
|
|
|
|
def test_all_definitive_rejections_exhaust_route(self):
|
|
result = provider_resolution.resolve_provider_key(
|
|
'sk-' + ('f' * 32),
|
|
hint='ambiguous_qwen_deepseek',
|
|
configured_order=['deepseek', 'qwen'],
|
|
probe=lambda *_args: {'status': 'DEAD', 'http_status': 401},
|
|
)
|
|
|
|
self.assertEqual(result['status'], 'DEAD')
|
|
self.assertEqual(result['provider_resolution'], 'exhausted')
|
|
self.assertEqual(len(result['provider_resolution_attempts']), 2)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|