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