sashatrask
239 lines
11 KiB
Python
239 lines
11 KiB
Python
import json
|
|
import os
|
|
from pathlib import Path
|
|
import sys
|
|
import tempfile
|
|
import unittest
|
|
from unittest import mock
|
|
|
|
|
|
ROOT = Path(__file__).resolve().parents[1]
|
|
APP_DIR = ROOT / 'app'
|
|
sys.path.insert(0, str(APP_DIR))
|
|
|
|
import supervisor
|
|
from lifecycle_authority import (
|
|
DISCOVERY_PRODUCER_ROLE,
|
|
DISCOVERY_PRODUCER_SOURCES,
|
|
)
|
|
from runtime_security import ensure_private_directory, harden_private_file
|
|
|
|
|
|
def producer_fixture(temp_dir, source='gitlab', source_config=None):
|
|
ensure_private_directory(temp_dir, reject_reparse=True)
|
|
config_path = os.path.join(temp_dir, 'config.yaml')
|
|
Path(config_path).write_text('{}\n', encoding='ascii')
|
|
harden_private_file(config_path)
|
|
log_dir = os.path.join(temp_dir, 'logs')
|
|
state_dir = os.path.join(temp_dir, 'state')
|
|
ensure_private_directory(log_dir, reject_reparse=True)
|
|
ensure_private_directory(state_dir, reject_reparse=True)
|
|
canonical_dsn = 'postgresql://managed@127.0.0.1/truf'
|
|
gate = supervisor.DependencyGate(ready=True, database_url=canonical_dsn)
|
|
child_environment = {
|
|
'TRUF_SUPERVISOR_TOKEN': 'trusted-token',
|
|
'TRUF_SUPERVISOR_CHILD_KIND': DISCOVERY_PRODUCER_ROLE,
|
|
'TRUF_MANAGED_POSTGRES_DSN': canonical_dsn,
|
|
}
|
|
producer = supervisor.ManagedDiscoveryProducer(
|
|
source, config_path, str(APP_DIR), temp_dir,
|
|
{
|
|
'log_dir': log_dir,
|
|
'state_dir': state_dir,
|
|
'defaults': {'enabled': True, 'restart': True},
|
|
'sources': {source: dict(source_config or {})},
|
|
},
|
|
{'enabled': True},
|
|
dependency_gate=gate,
|
|
child_environment=child_environment,
|
|
)
|
|
return producer, canonical_dsn
|
|
|
|
|
|
class DiscoveryProducerSupervisorTests(unittest.TestCase):
|
|
def test_allowlist_and_authenticated_one_shot_command_are_exact(self):
|
|
self.assertEqual(
|
|
DISCOVERY_PRODUCER_SOURCES,
|
|
('gitlab', 'dockerhub', 'huggingface'),
|
|
)
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
for source in DISCOVERY_PRODUCER_SOURCES:
|
|
with self.subTest(source=source):
|
|
producer, _ = producer_fixture(temp_dir, source=source)
|
|
command = producer.build_command()
|
|
self.assertEqual(command[:4], [sys.executable, '-I', '-S', '-B'])
|
|
self.assertEqual(Path(command[4]).name, 'child_bootstrap.py')
|
|
self.assertEqual(command[5:7], [DISCOVERY_PRODUCER_ROLE, '--'])
|
|
self.assertEqual(command[-5:], [
|
|
'--config', producer.config_path, '--source', source, '--once',
|
|
])
|
|
self.assertTrue(producer.once)
|
|
self.assertTrue(producer.repeat)
|
|
self.assertEqual(producer.producer_id, f'discovery-producer:{source}')
|
|
self.assertTrue(producer.log_path.endswith(
|
|
f'discovery-producer-{source}.log',
|
|
))
|
|
self.assertTrue(producer.state_path.endswith(
|
|
f'runner_state_{source}.json',
|
|
))
|
|
|
|
with self.assertRaisesRegex(ValueError, 'allowlist'):
|
|
producer_fixture(temp_dir, source='github')
|
|
|
|
def test_environment_strips_overrides_and_restores_immutable_authority(self):
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
producer, canonical_dsn = producer_fixture(temp_dir, source_config={
|
|
'enabled': True,
|
|
'env': {
|
|
'PGPASSWORD': 'forged',
|
|
'SCANNER_DB_URL': 'postgresql://forged',
|
|
'TRUF_SUPERVISOR_TOKEN': 'forged-token',
|
|
'TRUF_SUPERVISOR_CHILD_KIND': 'scanner',
|
|
'TRUF_MANAGED_POSTGRES_DSN': 'postgresql://forged',
|
|
'SAFE_PROVIDER_COMPAT': 'retained',
|
|
},
|
|
})
|
|
environment = producer.build_env()
|
|
self.assertNotIn('PGPASSWORD', environment)
|
|
self.assertEqual(environment['SCANNER_DB_URL'], canonical_dsn)
|
|
self.assertEqual(environment['DATABASE_URL'], canonical_dsn)
|
|
self.assertEqual(environment['TRUF_MANAGED_POSTGRES_DSN'], canonical_dsn)
|
|
self.assertEqual(environment['TRUF_SUPERVISOR_TOKEN'], 'trusted-token')
|
|
self.assertEqual(environment['TRUF_SUPERVISOR_CHILD_KIND'], DISCOVERY_PRODUCER_ROLE)
|
|
self.assertEqual(environment['TRUF_DISCOVERY_SOURCE'], 'gitlab')
|
|
self.assertEqual(environment['SAFE_PROVIDER_COMPAT'], 'retained')
|
|
self.assertNotIn('SCANNER_SKIP_STARTUP_CLEANUP', environment)
|
|
|
|
def test_reconfigure_cannot_disable_one_shot_repeat_mode(self):
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
producer, _ = producer_fixture(temp_dir)
|
|
producer.reconfigure(
|
|
temp_dir,
|
|
{
|
|
'log_dir': os.path.join(temp_dir, 'logs'),
|
|
'state_dir': os.path.join(temp_dir, 'state'),
|
|
'defaults': {'enabled': True},
|
|
},
|
|
{'enabled': True, 'once': False, 'repeat': False, 'extra_args': ['--loop']},
|
|
)
|
|
self.assertTrue(producer.once)
|
|
self.assertTrue(producer.repeat)
|
|
self.assertNotIn('--loop', producer.build_command())
|
|
|
|
def test_structured_state_is_bounded_and_signature_remains_nine_fields(self):
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
producers = []
|
|
for source in DISCOVERY_PRODUCER_SOURCES:
|
|
producer, _ = producer_fixture(temp_dir, source=source)
|
|
producers.append(producer)
|
|
state = {
|
|
'sources': {
|
|
source: {
|
|
'last_status': 'failed',
|
|
'last_cycle_result': {
|
|
'status': 'must-not-leak',
|
|
'fetched_count': 7,
|
|
'queued_new_count': 5,
|
|
'queued_updated_count': 2,
|
|
'provider_error': 'must-not-leak',
|
|
},
|
|
'last_discovery_success_at': 'must-not-leak',
|
|
'last_error_category': 'must-not-leak',
|
|
'raw_error': 'must-not-leak',
|
|
},
|
|
},
|
|
'token': 'must-not-leak',
|
|
}
|
|
Path(producer.state_path).write_text(json.dumps(state), encoding='utf-8')
|
|
snapshot = producer.structured_state()
|
|
serialized = json.dumps(snapshot, sort_keys=True)
|
|
self.assertEqual(snapshot['id'], f'discovery-producer:{source}')
|
|
self.assertEqual(snapshot['source'], source)
|
|
self.assertEqual(snapshot['role'], DISCOVERY_PRODUCER_ROLE)
|
|
self.assertEqual(snapshot['process_state'], 'stopped')
|
|
self.assertEqual(snapshot['interval_seconds'], producer.interval)
|
|
self.assertTrue(snapshot['restart_enabled'])
|
|
self.assertEqual(snapshot['allowed_actions'], [
|
|
'start', 'stop', 'restart', 'pause', 'resume', 'set-interval',
|
|
])
|
|
self.assertEqual(snapshot['safe_error_category'], 'runtime_error')
|
|
self.assertIsNone(snapshot['last_successful_discovery_at'])
|
|
self.assertEqual(snapshot['last_cycle_result'], {
|
|
'status': 'unknown',
|
|
'fetched_count': 7,
|
|
'queued_new_count': 5,
|
|
'queued_updated_count': 2,
|
|
})
|
|
self.assertNotIn('must-not-leak', serialized)
|
|
signature = supervisor.table_signature(producers)
|
|
self.assertEqual(len(signature), len(DISCOVERY_PRODUCER_SOURCES))
|
|
self.assertTrue(all(len(row) == 9 for row in signature))
|
|
|
|
def test_typed_restart_targets_only_one_exact_producer(self):
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
producers = [
|
|
producer_fixture(temp_dir, source=source)[0]
|
|
for source in DISCOVERY_PRODUCER_SOURCES
|
|
]
|
|
restart_methods = {}
|
|
for producer in producers:
|
|
replacement = mock.Mock(return_value=True)
|
|
restart_methods[producer.producer_id] = replacement
|
|
producer.restart_now = replacement
|
|
|
|
selected = producers[1]
|
|
result = supervisor.run_managed_source_action({
|
|
'schema': supervisor.CONTROL_SCHEMA,
|
|
'instance_id': 'fixture-instance',
|
|
'token': 't' * 48,
|
|
'action': 'managed-source-action',
|
|
'source_id': selected.producer_id,
|
|
'source_action': 'restart',
|
|
}, producers, {})
|
|
self.assertEqual(result['source']['id'], selected.producer_id)
|
|
restart_methods[selected.producer_id].assert_called_once_with()
|
|
for producer_id, method in restart_methods.items():
|
|
if producer_id != selected.producer_id:
|
|
method.assert_not_called()
|
|
|
|
for source_id in ('discovery-producer:github', 'discovery-producer:unknown'):
|
|
with self.subTest(source_id=source_id), self.assertRaisesRegex(
|
|
ValueError, 'unknown',
|
|
):
|
|
supervisor.run_managed_source_action({
|
|
'schema': supervisor.CONTROL_SCHEMA,
|
|
'instance_id': 'fixture-instance',
|
|
'token': 't' * 48,
|
|
'action': 'managed-source-action',
|
|
'source_id': source_id,
|
|
'source_action': 'restart',
|
|
}, producers, {})
|
|
self.assertEqual(
|
|
sum(method.call_count for method in restart_methods.values()), 1,
|
|
)
|
|
|
|
def test_restart_does_not_relaunch_after_start_gate_closes(self):
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
producer, _ = producer_fixture(temp_dir)
|
|
gate = {'open': True}
|
|
producer.start_gate = lambda: gate['open']
|
|
|
|
def stop_for_shutdown(timeout=10):
|
|
gate['open'] = False
|
|
return True
|
|
|
|
with (
|
|
mock.patch.object(producer, 'stop', side_effect=stop_for_shutdown),
|
|
mock.patch.object(supervisor, 'OwnedProcess') as launch,
|
|
):
|
|
self.assertFalse(producer.restart_now())
|
|
launch.assert_not_called()
|
|
self.assertEqual(
|
|
producer.last_action_error,
|
|
'supervisor lifecycle start gate is closed',
|
|
)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|