sashatrask
448 lines
19 KiB
Python
448 lines
19 KiB
Python
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()
|