Initial server source import
This commit is contained in:
sashatrask
@@ -0,0 +1,272 @@
|
||||
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()
|
||||
Reference in New Issue
Block a user