Initial server source import
This commit is contained in:
sashatrask
@@ -0,0 +1,447 @@
|
||||
import ast
|
||||
import inspect
|
||||
import os
|
||||
import sys
|
||||
import tempfile
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
from unittest import mock
|
||||
|
||||
|
||||
ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
APP = os.path.join(ROOT, 'app')
|
||||
if APP not in sys.path:
|
||||
sys.path.insert(0, APP)
|
||||
|
||||
import console_runner
|
||||
from scanner_db import DiscoveryPausedError
|
||||
|
||||
|
||||
class DiscoveryDB:
|
||||
def __init__(self, *, paused=False, enqueue_error=None):
|
||||
self.conn = SimpleNamespace(is_postgres=True)
|
||||
self.paused = paused
|
||||
self.enqueue_error = enqueue_error
|
||||
self.enqueues = []
|
||||
self.observations = []
|
||||
self.finished = []
|
||||
|
||||
def require_runtime_safety_schema(self):
|
||||
return True
|
||||
|
||||
def require_final_cutover(self):
|
||||
return True
|
||||
|
||||
def runtime_control_state(self):
|
||||
return {'effective_discovery_paused': self.paused}
|
||||
|
||||
def known_target_normalizations_for(self, source, platform, targets):
|
||||
return set()
|
||||
|
||||
def enqueue_targets(
|
||||
self, source, platform, query, targets, requeue_done=False,
|
||||
unresolved_targets=None, *, discovery_admission=False,
|
||||
):
|
||||
if self.enqueue_error is not None:
|
||||
raise self.enqueue_error
|
||||
record = {
|
||||
'source': source,
|
||||
'platform': platform,
|
||||
'query': query,
|
||||
'targets': list(targets),
|
||||
'unresolved_targets': list(unresolved_targets or []),
|
||||
'discovery_admission': discovery_admission,
|
||||
}
|
||||
self.enqueues.append(record)
|
||||
return len(record['targets']) + len(record['unresolved_targets'])
|
||||
|
||||
def observe_discovered_targets(
|
||||
self, source, platform, query, discoveries, *, rescan_limit,
|
||||
cooldown_seconds,
|
||||
):
|
||||
self.observations.append({
|
||||
'source': source,
|
||||
'platform': platform,
|
||||
'query': query,
|
||||
'discoveries': list(discoveries),
|
||||
'rescan_limit': rescan_limit,
|
||||
'cooldown_seconds': cooldown_seconds,
|
||||
})
|
||||
return {
|
||||
'attempted_count': len(discoveries),
|
||||
'queued_new_count': len(discoveries),
|
||||
'queued_updated_count': 0,
|
||||
}
|
||||
|
||||
def finish_source_cycle(self, *args):
|
||||
self.finished.append(args)
|
||||
return True
|
||||
|
||||
def has_claimable_targets(self, *args, **kwargs):
|
||||
raise AssertionError('discovery-only cycle probed scanner backlog')
|
||||
|
||||
def has_claimable_targets_v2(self, *args, **kwargs):
|
||||
raise AssertionError('discovery-only cycle probed scanner backlog')
|
||||
|
||||
|
||||
def args_for(source):
|
||||
platforms = {
|
||||
'gitlab': 'gitlab',
|
||||
'dockerhub': 'docker',
|
||||
'huggingface': 'huggingface',
|
||||
}
|
||||
return SimpleNamespace(
|
||||
platform=platforms[source],
|
||||
query='credential',
|
||||
updated_target_rescan_enabled=False,
|
||||
)
|
||||
|
||||
|
||||
class DiscoveryOnlyCycleTests(unittest.TestCase):
|
||||
def test_discovery_producer_mode_uses_only_discovery_runner_and_postgres(self):
|
||||
class ProducerDB:
|
||||
enabled = True
|
||||
conn = SimpleNamespace(is_postgres=True)
|
||||
|
||||
def __init__(self):
|
||||
self.application_name = None
|
||||
self.started = None
|
||||
self.finished = []
|
||||
self.closed = False
|
||||
|
||||
def set_application_name(self, value):
|
||||
self.application_name = value
|
||||
|
||||
def require_runtime_safety_schema(self):
|
||||
return True
|
||||
|
||||
def require_final_cutover(self):
|
||||
return True
|
||||
|
||||
def start_run(self, *args, **kwargs):
|
||||
self.started = (args, kwargs)
|
||||
return 41
|
||||
|
||||
def finish_run(self, *args):
|
||||
self.finished.append(args)
|
||||
|
||||
def close(self):
|
||||
self.closed = True
|
||||
|
||||
database = ProducerDB()
|
||||
state = {'sources': {'gitlab': console_runner.default_source_state()}}
|
||||
config = {
|
||||
'global': {'database_url': 'postgresql://managed'},
|
||||
'sources': {'gitlab': {'enabled': True}},
|
||||
}
|
||||
args = SimpleNamespace(
|
||||
config='config.yaml', source='gitlab', once=True,
|
||||
show_state=False, cleanup_only=False,
|
||||
)
|
||||
with mock.patch.object(console_runner, 'preflight_lifecycle_paths') as preflight, \
|
||||
mock.patch.object(console_runner, 'apply_global_config'), \
|
||||
mock.patch.object(console_runner, 'require_sensitive_runtime_paths'), \
|
||||
mock.patch.object(console_runner, 'get_state_path', return_value='state.json'), \
|
||||
mock.patch.object(console_runner, 'load_state', return_value=state), \
|
||||
mock.patch.object(console_runner, 'load_secrets', return_value={}), \
|
||||
mock.patch.object(console_runner, 'refresh_auth_summary'), \
|
||||
mock.patch.object(console_runner, 'save_state'), \
|
||||
mock.patch.object(console_runner, 'hash_file', return_value='a' * 64), \
|
||||
mock.patch.object(console_runner, 'ScannerDB', return_value=database), \
|
||||
mock.patch.object(console_runner, 'run_configured_source', return_value={}) as configured, \
|
||||
mock.patch.object(console_runner, 'initialize_scanner_runtime') as scanner_init, \
|
||||
mock.patch.object(console_runner, 'check_dependencies') as dependencies, \
|
||||
mock.patch.object(console_runner, 'cleanup_pending_command_work_dirs') as cleanup_work, \
|
||||
mock.patch.object(console_runner, 'cleanup_stale_temp_dirs') as cleanup_temp:
|
||||
console_runner.run_discovery_producer_mode(args, config)
|
||||
|
||||
preflight.assert_called_once_with(
|
||||
'config.yaml', config, authority_profile='discovery-producer',
|
||||
)
|
||||
scanner_init.assert_not_called()
|
||||
dependencies.assert_not_called()
|
||||
cleanup_work.assert_not_called()
|
||||
cleanup_temp.assert_not_called()
|
||||
self.assertEqual(database.application_name, 'truf-discovery:gitlab')
|
||||
self.assertEqual(database.started[0][0], 'discovery-producer')
|
||||
self.assertEqual(database.started[1]['selected_source'], 'gitlab')
|
||||
self.assertEqual(database.started[1]['selected_platform'], 'gitlab')
|
||||
self.assertTrue(database.closed)
|
||||
configured.assert_called_once()
|
||||
self.assertIs(configured.call_args.kwargs['cycle_runner'], console_runner.run_discovery_cycle)
|
||||
|
||||
def test_discovery_producer_mode_requires_exact_one_shot_core_source(self):
|
||||
config = {'global': {}, 'sources': {
|
||||
'gitlab': {}, 'dockerhub': {}, 'huggingface': {}, 'github': {}, 'docker': {},
|
||||
}}
|
||||
cases = (
|
||||
(False, 'gitlab'),
|
||||
(True, None),
|
||||
(True, 'github'),
|
||||
(True, 'docker'),
|
||||
)
|
||||
for once, source in cases:
|
||||
args = SimpleNamespace(
|
||||
config='config.yaml', source=source, once=once,
|
||||
show_state=False, cleanup_only=False,
|
||||
)
|
||||
with self.subTest(once=once, source=source), \
|
||||
mock.patch.object(console_runner, 'preflight_lifecycle_paths') as preflight, \
|
||||
self.assertRaises(SystemExit):
|
||||
console_runner.run_discovery_producer_mode(args, config)
|
||||
preflight.assert_not_called()
|
||||
|
||||
def run_cycle(self, source, fetched, *, database=None):
|
||||
database = database or DiscoveryDB()
|
||||
with mock.patch.object(
|
||||
console_runner, 'fetch_targets', return_value=fetched,
|
||||
) as fetch:
|
||||
metrics = console_runner.run_discovery_cycle(
|
||||
args_for(source), database, 11, 17, source,
|
||||
)
|
||||
fetch.assert_called_once()
|
||||
return database, metrics
|
||||
|
||||
def assert_zero_scan_metrics(self, metrics):
|
||||
self.assertEqual(metrics['scan_requested_count'], 0)
|
||||
self.assertEqual(metrics['scanned_count'], 0)
|
||||
self.assertEqual(metrics['staged_count'], 0)
|
||||
self.assertEqual(metrics['findings_count'], 0)
|
||||
self.assertFalse(metrics['backlog_only'])
|
||||
|
||||
def test_gitlab_discovery_normalizes_and_uses_provider_admission(self):
|
||||
database, metrics = self.run_cycle(
|
||||
'gitlab',
|
||||
['https://gitlab.com/acme/one.git', 'https://gitlab.com/acme/two'],
|
||||
)
|
||||
self.assertEqual(metrics['cycle_status'], 'completed')
|
||||
self.assertEqual(metrics['fetched_count'], 2)
|
||||
self.assertEqual(metrics['queued_new_count'], 2)
|
||||
self.assertEqual(len(database.enqueues), 1)
|
||||
self.assertTrue(database.enqueues[0]['discovery_admission'])
|
||||
self.assertEqual(database.enqueues[0]['source'], 'gitlab')
|
||||
self.assertEqual(len(database.finished), 1)
|
||||
self.assert_zero_scan_metrics(metrics)
|
||||
|
||||
def test_huggingface_discovery_preserves_newest_remote_metadata(self):
|
||||
database = DiscoveryDB()
|
||||
args = args_for('huggingface')
|
||||
args.updated_target_rescan_enabled = True
|
||||
args.updated_target_rescan_max_per_cycle = 4
|
||||
args.updated_target_rescan_cooldown_hours = 12
|
||||
fetched = [
|
||||
{'target': 'acme/space', 'remote_modified_at': '2026-09-18T00:00:00Z'},
|
||||
{'target': 'acme/space', 'remote_modified_at': '2026-09-19T00:00:00Z'},
|
||||
]
|
||||
with mock.patch.object(
|
||||
console_runner, 'fetch_targets', return_value=fetched,
|
||||
):
|
||||
metrics = console_runner.run_discovery_cycle(
|
||||
args, database, 11, 17, 'huggingface',
|
||||
)
|
||||
self.assertEqual(database.observations[0]['discoveries'], [{
|
||||
'target': 'acme/space',
|
||||
'remote_modified_at': '2026-09-19T00:00:00+00:00',
|
||||
}])
|
||||
self.assertEqual(database.observations[0]['rescan_limit'], 4)
|
||||
self.assertEqual(database.observations[0]['cooldown_seconds'], 12 * 3600)
|
||||
self.assertEqual(metrics['queued_new_count'], 1)
|
||||
self.assert_zero_scan_metrics(metrics)
|
||||
|
||||
def test_huggingface_discovery_filters_explicit_restrictions(self):
|
||||
database = DiscoveryDB()
|
||||
args = args_for('huggingface')
|
||||
args.updated_target_rescan_enabled = True
|
||||
args.updated_target_rescan_max_per_cycle = 4
|
||||
args.updated_target_rescan_cooldown_hours = 12
|
||||
records = [
|
||||
{'target': 'acme/public', 'private': False, 'gated': None},
|
||||
{'target': 'acme/metadata-omitted'},
|
||||
{'target': 'acme/private', 'private': True},
|
||||
{'target': 'acme/protected', 'protected': True},
|
||||
{'target': 'acme/gated', 'gated': 'manual'},
|
||||
{'target': 'acme/disabled', 'disabled': True},
|
||||
{'target': 'acme/malformed', 'private': 'false'},
|
||||
]
|
||||
projected, _metrics = console_runner.enqueue_discovered_targets(
|
||||
args, records, database, 11, 17, 'huggingface',
|
||||
)
|
||||
self.assertEqual(projected, ['acme/public', 'acme/metadata-omitted'])
|
||||
self.assertEqual(
|
||||
[item['target'] for item in database.observations[0]['discoveries']],
|
||||
['acme/public', 'acme/metadata-omitted'],
|
||||
)
|
||||
|
||||
def test_huggingface_fetch_uses_existing_metadata_without_second_request(self):
|
||||
args = SimpleNamespace(
|
||||
platform='huggingface', mode='search', token='discovery-token',
|
||||
pages=1, fetch_timeout=15, updated_target_rescan_enabled=False,
|
||||
stop_on_seen_pages=False, huggingface_discovery_request_attempts=1,
|
||||
huggingface_discovery_retry_delay=0,
|
||||
)
|
||||
spaces = [
|
||||
{'url': 'acme/public', 'updated_at': '2026-09-19T00:00:00Z', 'private': False},
|
||||
{'url': 'acme/private', 'updated_at': '2026-09-19T00:00:00Z', 'private': True},
|
||||
{'url': 'acme/protected', 'protected': True},
|
||||
{'url': 'acme/gated', 'gated': 'manual'},
|
||||
{'url': 'acme/disabled', 'disabled': True},
|
||||
{'url': 'acme/metadata-omitted'},
|
||||
]
|
||||
with mock.patch.object(
|
||||
console_runner, 'fetch_huggingface_spaces', return_value=spaces,
|
||||
) as fetch:
|
||||
targets = console_runner.fetch_targets(args)
|
||||
self.assertEqual(
|
||||
[item['target'] for item in targets],
|
||||
['acme/public', 'acme/metadata-omitted'],
|
||||
)
|
||||
self.assertFalse(targets[0]['private'])
|
||||
self.assertEqual(targets[0]['remote_modified_at'], '2026-09-19T00:00:00+00:00')
|
||||
self.assertTrue(fetch.call_args.kwargs['return_metadata'])
|
||||
|
||||
def test_dockerhub_discovery_enqueues_bare_repositories_as_unresolved(self):
|
||||
digest = 'library/busybox@sha256:' + ('a' * 64)
|
||||
database, metrics = self.run_cycle(
|
||||
'dockerhub', ['library/alpine', digest],
|
||||
)
|
||||
admission = database.enqueues[0]
|
||||
self.assertEqual(admission['targets'], [digest])
|
||||
self.assertEqual(admission['unresolved_targets'], ['library/alpine'])
|
||||
self.assertTrue(admission['discovery_admission'])
|
||||
self.assertEqual(metrics['queued_new_count'], 2)
|
||||
self.assert_zero_scan_metrics(metrics)
|
||||
|
||||
def test_every_core_producer_runs_with_backlog_and_zero_scan_work(self):
|
||||
discoveries = {
|
||||
'gitlab': ['https://gitlab.com/acme/project'],
|
||||
'dockerhub': ['library/alpine'],
|
||||
'huggingface': ['acme/space'],
|
||||
}
|
||||
for source, fetched in discoveries.items():
|
||||
with self.subTest(source=source), \
|
||||
mock.patch.object(console_runner, 'prepare_targets') as prepare, \
|
||||
mock.patch.object(console_runner, 'run_cycle_v2') as scanner_cycle, \
|
||||
mock.patch.object(console_runner, 'scan_target_result') as scan, \
|
||||
mock.patch.object(console_runner, 'stage_result_bundle') as stage:
|
||||
database, metrics = self.run_cycle(source, fetched)
|
||||
prepare.assert_not_called()
|
||||
scanner_cycle.assert_not_called()
|
||||
scan.assert_not_called()
|
||||
stage.assert_not_called()
|
||||
self.assertEqual(metrics['fetched_count'], 1)
|
||||
self.assertEqual(metrics['cycle_status'], 'completed')
|
||||
self.assertEqual(len(database.finished), 1)
|
||||
self.assert_zero_scan_metrics(metrics)
|
||||
|
||||
def test_persistent_pause_skips_provider_and_admission_for_every_core_source(self):
|
||||
for source in ('gitlab', 'dockerhub', 'huggingface'):
|
||||
with self.subTest(source=source):
|
||||
database = DiscoveryDB(paused=True)
|
||||
with mock.patch.object(console_runner, 'fetch_targets') as fetch:
|
||||
metrics = console_runner.run_discovery_cycle(
|
||||
args_for(source), database, 11, 17, source,
|
||||
)
|
||||
fetch.assert_not_called()
|
||||
self.assertEqual(database.enqueues, [])
|
||||
self.assertEqual(metrics['cycle_status'], 'paused')
|
||||
self.assertEqual(database.finished[0][1], 'paused')
|
||||
self.assert_zero_scan_metrics(metrics)
|
||||
|
||||
def test_pause_racing_admission_finishes_once_as_paused(self):
|
||||
control = {'effective_discovery_paused': True}
|
||||
database = DiscoveryDB(enqueue_error=DiscoveryPausedError(control))
|
||||
with mock.patch.object(
|
||||
console_runner, 'fetch_targets', return_value=['https://gitlab.com/acme/one'],
|
||||
):
|
||||
metrics = console_runner.run_discovery_cycle(
|
||||
args_for('gitlab'), database, 11, 17, 'gitlab',
|
||||
)
|
||||
self.assertEqual(metrics['cycle_status'], 'paused')
|
||||
self.assertEqual(metrics['queued_new_count'], 0)
|
||||
self.assertEqual(len(database.finished), 1)
|
||||
self.assert_zero_scan_metrics(metrics)
|
||||
|
||||
def test_empty_discovery_finishes_once_with_zero_scan_work(self):
|
||||
database, metrics = self.run_cycle('huggingface', [])
|
||||
self.assertEqual(len(database.finished), 1)
|
||||
self.assertEqual(database.finished[0][1], 'completed')
|
||||
self.assert_zero_scan_metrics(metrics)
|
||||
|
||||
def test_discovery_call_graph_has_no_scanner_claim_or_bundle_operations(self):
|
||||
forbidden = {
|
||||
'run_cycle', 'run_cycle_v2', '_run_cycle_legacy_compat',
|
||||
'prepare_targets', 'prepare_scan_options', 'acquire_scan_slot',
|
||||
'acquire_scan_slot_leases', 'scan_slot_scope', 'claim_targets',
|
||||
'reserve_and_claim_target', 'recover_result_reservation_claim',
|
||||
'has_claimable_targets', 'has_claimable_targets_v2',
|
||||
'resolve_due_docker_queue_targets_if_scan_queue_empty',
|
||||
'resolve_and_bind_git_claim',
|
||||
'resolve_and_bind_docker_claim', 'scan_target_result', 'scan_targets_batch',
|
||||
'result_spool_for_args', 'reserve_result_spool_claims',
|
||||
'prepare_scan_event', 'stage_scan_result_in_scope', 'stage_result_bundle',
|
||||
'ensure_bundle_reservation_paths', 'mark_result_bundle_ready',
|
||||
}
|
||||
names = set()
|
||||
for function in (
|
||||
console_runner.run_discovery_producer_mode,
|
||||
console_runner.run_discovery_cycle,
|
||||
console_runner.enqueue_discovered_targets,
|
||||
console_runner.run_dockerhub_incremental_discovery,
|
||||
console_runner.process_dockerhub_discovery_retry,
|
||||
console_runner.resolve_due_docker_queue_targets,
|
||||
console_runner.resolve_due_docker_experiment_targets,
|
||||
):
|
||||
tree = ast.parse(inspect.getsource(function))
|
||||
for node in ast.walk(tree):
|
||||
if isinstance(node, ast.Call):
|
||||
target = node.func
|
||||
if isinstance(target, ast.Name):
|
||||
names.add(target.id)
|
||||
elif isinstance(target, ast.Attribute):
|
||||
names.add(target.attr)
|
||||
self.assertEqual(names & forbidden, set())
|
||||
|
||||
def test_configured_paused_cycle_does_not_advance_query(self):
|
||||
class ConfiguredDB:
|
||||
conn = SimpleNamespace(is_postgres=True)
|
||||
|
||||
def start_source_cycle(self, *args, **kwargs):
|
||||
return 17
|
||||
|
||||
config = {
|
||||
'global': {},
|
||||
'sources': {
|
||||
'gitlab': {'queries': ['first', 'second'], 'mode': 'recent'},
|
||||
},
|
||||
}
|
||||
state = {'sources': {'gitlab': console_runner.default_source_state()}}
|
||||
args = SimpleNamespace(platform='gitlab', mode='recent')
|
||||
runner = mock.Mock(return_value={
|
||||
'cycle_status': 'paused',
|
||||
'scan_requested_count': 0,
|
||||
'staged_count': 0,
|
||||
'scanned_count': 0,
|
||||
'source_failure_count': 0,
|
||||
})
|
||||
with tempfile.TemporaryDirectory() as temp_dir, \
|
||||
mock.patch.object(console_runner, 'save_state'), \
|
||||
mock.patch.object(console_runner, 'select_auth_entry', return_value=None), \
|
||||
mock.patch.object(console_runner, 'refresh_auth_summary'), \
|
||||
mock.patch.object(
|
||||
console_runner, 'build_args_from_source_config', return_value=args,
|
||||
), \
|
||||
mock.patch.object(console_runner, 'configure_source_auth'), \
|
||||
mock.patch.object(console_runner, 'queue_counts_for_args', return_value={}):
|
||||
metrics = console_runner.run_configured_source(
|
||||
'gitlab', config, state, os.path.join(temp_dir, 'state.json'), {},
|
||||
ConfiguredDB(), run_id=3, cycle_runner=runner,
|
||||
)
|
||||
self.assertEqual(metrics['cycle_status'], 'paused')
|
||||
self.assertEqual(state['sources']['gitlab']['query_index'], 0)
|
||||
self.assertEqual(state['sources']['gitlab']['last_status'], 'paused')
|
||||
self.assertEqual(state['sources']['gitlab']['last_scanned'], 0)
|
||||
runner.assert_called_once()
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user