sashatrask
273 lines
10 KiB
Python
273 lines
10 KiB
Python
import os
|
|
from pathlib import Path
|
|
import re
|
|
import sys
|
|
from types import SimpleNamespace
|
|
import unittest
|
|
from unittest import mock
|
|
|
|
|
|
ROOT = Path(__file__).resolve().parents[1]
|
|
APP_DIR = ROOT / 'app'
|
|
sys.path.insert(0, str(APP_DIR))
|
|
|
|
import dashboard
|
|
import keycheck_runner
|
|
import scanner
|
|
from keycheck_candidates import extract_candidates, extract_structured_candidates
|
|
from keycheckers.kimi import kimiKeycheck as kimi
|
|
|
|
|
|
class KimiDetectorTests(unittest.TestCase):
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
import yaml
|
|
|
|
policy = yaml.safe_load(
|
|
(APP_DIR / 'trufflehog-custom-detectors.yaml').read_text(encoding='utf-8')
|
|
)
|
|
cls.detector = next(
|
|
item for item in policy['detectors'] if item['name'] == 'KimiMoonshot'
|
|
)
|
|
|
|
def test_detector_requires_explicit_kimi_or_moonshot_assignment(self):
|
|
self.assertEqual(
|
|
set(self.detector['keywords']), {'MOONSHOT_API_KEY', 'KIMI_API_KEY'},
|
|
)
|
|
pattern = re.compile(next(iter(self.detector['regex'].values())))
|
|
key = 'sk-' + ('K' * 40)
|
|
for text in (
|
|
f'MOONSHOT_API_KEY={key}',
|
|
f'"KIMI_API_KEY": "{key}"',
|
|
):
|
|
with self.subTest(text=text):
|
|
self.assertEqual(pattern.search(text).group(1), key)
|
|
self.assertIsNone(pattern.search(f'OPENAI_API_KEY={key}\nmodel=kimi-k2'))
|
|
self.assertIsNone(pattern.search(f'base_url=https://api.moonshot.ai/v1\nkey={key}'))
|
|
|
|
def test_detector_and_structured_context_create_kimi_candidates(self):
|
|
key = 'sk-' + ('A' * 40)
|
|
direct = next(iter(extract_candidates({
|
|
'DetectorName': 'KimiMoonshot', 'Raw': key,
|
|
})))
|
|
structured = next(iter(extract_structured_candidates({
|
|
'origin': 'fixture:postman',
|
|
'contexts': [{
|
|
'key': 'MOONSHOT_API_KEY', 'value': key, 'path': '$.moonshot',
|
|
}],
|
|
})))
|
|
self.assertEqual(direct.service, 'kimi')
|
|
self.assertEqual(structured.service, 'kimi')
|
|
self.assertEqual(direct.provider_key_hash, structured.provider_key_hash)
|
|
|
|
def test_explicit_kimi_evidence_routes_duplicate_generic_findings_only_to_kimi(self):
|
|
key = 'sk-' + ('B' * 40)
|
|
findings = [
|
|
{
|
|
'DetectorName': 'CustomRegex',
|
|
'ExtraData': {'name': 'KimiMoonshot'},
|
|
'Raw': key,
|
|
},
|
|
{'DetectorName': 'Qwen', 'Raw': key},
|
|
{'DetectorName': 'DeepSeek', 'Raw': key},
|
|
]
|
|
scanner.strip_nearby_context_for_persistence({'findings': findings})
|
|
self.assertEqual(
|
|
[item['ScannerContext']['provider_hint'] for item in findings],
|
|
['kimi', 'kimi', 'kimi'],
|
|
)
|
|
|
|
|
|
class KimiKeycheckTests(unittest.TestCase):
|
|
KEY = 'sk-' + ('C' * 40)
|
|
|
|
@staticmethod
|
|
def response(status_code, payload, text=''):
|
|
response = mock.Mock(status_code=status_code, text=text)
|
|
response.json.return_value = payload
|
|
return response
|
|
|
|
def test_positive_balance_is_valid(self):
|
|
response = self.response(200, {
|
|
'code': 0,
|
|
'status': True,
|
|
'data': {
|
|
'available_balance': 12.5,
|
|
'voucher_balance': 2.5,
|
|
'cash_balance': 10,
|
|
},
|
|
})
|
|
with mock.patch.object(kimi.requests, 'get', return_value=response) as request:
|
|
result = kimi.check_base_url(
|
|
self.KEY, 'https://api.moonshot.ai/v1', None, 1,
|
|
)
|
|
self.assertEqual(result['status'], 'VALID')
|
|
self.assertEqual(result['balance_usd'], 12.5)
|
|
self.assertTrue(result['authenticated'])
|
|
self.assertEqual(
|
|
request.call_args.args[0],
|
|
'https://api.moonshot.ai/v1/users/me/balance',
|
|
)
|
|
|
|
def test_zero_balance_is_authenticated_but_not_usable(self):
|
|
response = self.response(200, {
|
|
'code': 0,
|
|
'status': True,
|
|
'data': {
|
|
'available_balance': 0,
|
|
'voucher_balance': 0,
|
|
'cash_balance': 0,
|
|
},
|
|
})
|
|
with mock.patch.object(kimi.requests, 'get', return_value=response):
|
|
result = kimi.check_base_url(
|
|
self.KEY, 'https://api.moonshot.ai/v1', None, 1,
|
|
)
|
|
self.assertEqual(result['status'], 'NO_BALANCE')
|
|
self.assertTrue(result['authenticated'])
|
|
|
|
def test_region_mismatch_does_not_hide_retryable_network_result(self):
|
|
attempts = [
|
|
{'status': 'DEAD', 'region': 'api.moonshot.ai'},
|
|
{'status': 'NETWORK', 'region': 'api.moonshot.cn', 'message': 'timeout'},
|
|
]
|
|
with mock.patch.object(kimi, 'check_base_url', side_effect=attempts):
|
|
result = kimi.check_key(
|
|
self.KEY, list(kimi.DEFAULT_BASE_URLS), None, 1,
|
|
)
|
|
self.assertEqual(result['status'], 'NETWORK')
|
|
|
|
def test_foreign_provider_prefix_is_rejected_without_request(self):
|
|
key = 'sk-ant-' + ('D' * 40)
|
|
with mock.patch.object(kimi.requests, 'get') as request:
|
|
result = kimi.check_key(key, list(kimi.DEFAULT_BASE_URLS), None, 1)
|
|
self.assertEqual(result['status'], 'UNKNOWN')
|
|
self.assertTrue(result['candidate_rejected'])
|
|
request.assert_not_called()
|
|
|
|
def test_explicit_candidate_uses_generic_sk_routing_gate(self):
|
|
finding = {
|
|
'DetectorName': 'KimiMoonshot',
|
|
'Raw': self.KEY,
|
|
'ScannerContext': {
|
|
'provider_hint': 'kimi',
|
|
'provider_hint_source': 'explicit_assignment',
|
|
},
|
|
}
|
|
item = {'raw': self.KEY, 'source': 'fixture', 'finding': finding}
|
|
with mock.patch.object(kimi, 'iter_findings', return_value=[item]), \
|
|
mock.patch.object(kimi, 'combined_provider_routing_hint', return_value='kimi') as routing, \
|
|
mock.patch.object(kimi, 'provider_routing_database_failed', return_value=False):
|
|
candidates = list(kimi.iter_candidate_keys('unused.jsonl', []))
|
|
self.assertEqual([candidate for candidate, _, _ in candidates], [self.KEY])
|
|
routing.assert_called_once_with(self.KEY, 'kimi')
|
|
|
|
def test_recheck_uses_authenticated_candidate_route(self):
|
|
finding = {'DetectorName': 'DeepSeek', 'Raw': self.KEY}
|
|
item = {
|
|
'raw': self.KEY,
|
|
'source': 'postgres-fixture',
|
|
'finding': finding,
|
|
'candidate_metadata': {'provider_hint': 'kimi'},
|
|
}
|
|
args = SimpleNamespace(
|
|
input='unused.jsonl', plain=[], proxy_file='', timeout=1, max_keys=1,
|
|
base_url=[], no_default_base_urls=False, retry_network=False,
|
|
retry_limited=False, retry_unknown=False, retry_restricted=False,
|
|
retry_no_balance=False, retry_valid=True, recheck_all=False,
|
|
debug=False,
|
|
)
|
|
result = {'status': 'DEAD', 'message': 'fixture invalid key'}
|
|
with mock.patch.dict(os.environ, {
|
|
'KEYCHECK_INPUT_MODE': 'postgres', 'KEYCHECK_SERVICE': 'kimi',
|
|
}, clear=False), mock.patch.object(
|
|
kimi, 'require_provider_authority',
|
|
), mock.patch.object(
|
|
kimi, 'parse_args', return_value=args,
|
|
), mock.patch.object(
|
|
kimi, 'ensure_files',
|
|
), mock.patch.object(
|
|
kimi, 'load_proxies', return_value=None,
|
|
), mock.patch.object(
|
|
kimi, 'load_checked_statuses', return_value={},
|
|
), mock.patch.object(
|
|
kimi, 'load_known_keys', return_value=set(),
|
|
), mock.patch.object(
|
|
kimi, 'iter_findings', return_value=iter([item]),
|
|
), mock.patch.object(
|
|
kimi, 'read_plain_keys', return_value=[],
|
|
), mock.patch.object(
|
|
kimi, 'should_skip_key', return_value=False,
|
|
), mock.patch.object(
|
|
kimi, 'base_urls_from_args', return_value=['https://fixture.invalid/v1'],
|
|
), mock.patch.object(
|
|
kimi, 'combined_provider_routing_hint', return_value='deepseek',
|
|
) as routing, mock.patch.object(
|
|
kimi, 'check_key', return_value=result,
|
|
) as check_key, mock.patch.object(
|
|
kimi, 'resolve_provider_key', return_value=result,
|
|
) as resolve, mock.patch.object(
|
|
kimi, 'write_result',
|
|
) as write_result:
|
|
kimi.main()
|
|
|
|
routing.assert_not_called()
|
|
check_key.assert_called_once()
|
|
resolve.assert_not_called()
|
|
self.assertNotIn(
|
|
kimi.CANDIDATE_PROVIDER_ROUTE_FIELD,
|
|
write_result.call_args.args[3],
|
|
)
|
|
|
|
|
|
class KimiRegistryTests(unittest.TestCase):
|
|
def test_kimi_is_registered_for_runtime_and_dashboard(self):
|
|
self.assertIn('kimi', keycheck_runner.SERVICES)
|
|
self.assertIn('kimi', keycheck_runner.SERVICE_CAPABILITIES)
|
|
self.assertIn('kimi', keycheck_runner.POSTGRES_STATUS_FILE_NAMES)
|
|
self.assertIn('kimi', dashboard.VALIDATION_SERVICES)
|
|
|
|
def test_retry_unknown_requeues_no_context_in_postgres(self):
|
|
class Database:
|
|
enabled = True
|
|
|
|
def __init__(self):
|
|
self.groups = None
|
|
|
|
@staticmethod
|
|
def set_application_name(_name):
|
|
return None
|
|
|
|
@staticmethod
|
|
def require_runtime_safety_schema():
|
|
return None
|
|
|
|
@staticmethod
|
|
def require_final_cutover():
|
|
return None
|
|
|
|
def enqueue_keycheck_rechecks(self, _service, groups, **_kwargs):
|
|
self.groups = groups
|
|
return 1
|
|
|
|
@staticmethod
|
|
def close():
|
|
return None
|
|
|
|
database = Database()
|
|
args = SimpleNamespace(
|
|
input_mode='postgres', retry_network=False, retry_limited=False,
|
|
retry_unknown=True, retry_restricted=False, retry_no_balance=False,
|
|
retry_valid=False, recheck_all=False, max_keys=0,
|
|
)
|
|
with mock.patch.object(keycheck_runner, 'ScannerDB', return_value=database):
|
|
total = keycheck_runner.enqueue_requested_postgres_rechecks(
|
|
{'database_url': 'postgresql://fixture'}, ['qwen'], args,
|
|
)
|
|
self.assertEqual(total, 1)
|
|
self.assertEqual(database.groups, {'unknown', 'no_context'})
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|