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()