import os from pathlib import Path import runpy from types import SimpleNamespace 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 console_runner import db_backend import lifecycle_authority import supervisor import uvicorn import worker_api import worker_assignment import managed_files from runtime_security import ensure_private_directory class ConfiguredWorkerServiceTests(unittest.TestCase): def _config(self, bundle_root, **worker_overrides): worker = { 'enabled': True, 'address': '127.0.0.1', 'port': 8766, 'sources': ['gitlab'], 'auth_entries': {'gitlab': 'selected'}, 'compatibility_profiles': {'windows-x86_64': {'fixture': True}}, 'assignment_ttl_seconds': 86400, 'assignment_ttl_seconds_by_source': {}, 'max_bundle_bytes': 64 * 1024 * 1024, 'reaper_interval_seconds': 60, 'reaper_batch_size': 1000, 'limit_concurrency': 64, } worker.update(worker_overrides) return { 'global': { 'database_url': 'postgresql://fixture', 'result_bundle_dir': bundle_root, 'result_bundle_max_event_bytes': 64 * 1024 * 1024, 'remote_assignment_reserve_bytes': 2 * 1024 * 1024, 'remote_assignment_max_active': 50, 'result_bundle_max_total_bytes': 512 * 1024 * 1024, 'projection_backlog_max_bytes': 256 * 1024 * 1024, 'projection_backlog_headroom_bytes': 128 * 1024 * 1024, 'keycheck_queue_max_items': 131072, 'keycheck_queue_max_bytes': 128 * 1024 * 1024, 'keycheck_candidates_per_event': 2000, 'keycheck_candidate_bytes_per_event': 2 * 1024 * 1024, }, 'supervisor': {'worker_api': worker}, 'sources': {'gitlab': {'enabled': True}}, } def _patch_config(self, config): args = SimpleNamespace( platform='gitlab', exact_git_planning_enabled=True, timeout=1800, ) return ( mock.patch.object(console_runner, 'load_config', return_value=config), mock.patch.object(console_runner, 'apply_global_config'), mock.patch.object(console_runner, 'load_secrets', return_value={'fixture': True}), mock.patch.object(console_runner, 'auth_pool_entries', return_value=([ {'name': 'other', 'token': 'not-selected'}, {'name': 'selected', 'token': 'selected-secret'}, ], None)), mock.patch.object(console_runner, 'build_args_from_source_config', return_value=args), mock.patch.object(db_backend, 'database_url_from_env', return_value='postgresql://fixture'), mock.patch.object(worker_assignment, 'RemoteAssignmentBuilder'), ) def test_factory_is_disabled_by_default(self): config = {'global': {}, 'supervisor': {}, 'sources': {}} with mock.patch.object(console_runner, 'load_config', return_value=config): with self.assertRaisesRegex(ValueError, 'not explicitly enabled'): worker_api.build_configured_worker_service('config.yaml', {'instance_id': 'fixture'}) def test_factory_rejects_public_or_unspecified_bind_before_database_access(self): config = self._config('unused', address='0.0.0.0') with mock.patch.object(console_runner, 'load_config', return_value=config), \ mock.patch.object(db_backend, 'database_url_from_env') as database_url: with self.assertRaisesRegex(ValueError, 'loopback or private'): worker_api.build_configured_worker_service( 'config.yaml', {'instance_id': 'fixture'}, ) database_url.assert_not_called() def test_factory_uses_canonical_dsn_and_explicit_auth_entry(self): with tempfile.TemporaryDirectory() as temp_dir: bundle_root = ensure_private_directory( os.path.join(temp_dir, 'bundles'), reject_reparse=True, ) config = self._config( bundle_root, assignment_ttl_seconds_by_source={'gitlab': 3660}, ) patches = self._patch_config(config) with patches[0] as load_config, patches[1] as apply_global, \ patches[2] as load_secrets, patches[3], patches[4] as build_args, \ patches[5], patches[6] as builder: service, runtime = worker_api.build_configured_worker_service( 'config.yaml', {'instance_id': 'supervisor-fixture'}, db_factory=object, ) load_config.assert_called_once_with( 'config.yaml', managed_postgres=True, final_cutover=True, ) apply_global.assert_called_once_with(config['global']) load_secrets.assert_called_once_with(config, 'config.yaml') self.assertEqual( build_args.call_args.kwargs['auth_entry']['token'], 'selected-secret', ) builder.assert_called_once_with( 'postgresql://fixture', bundle_root, {'gitlab': build_args.return_value}, config['supervisor']['worker_api']['compatibility_profiles'], 'supervisor-fixture', assignment_ttl_seconds=86400, assignment_ttl_seconds_by_source={'gitlab': 3660}, result_upload_body_timeout_seconds=1800, credential_refs={'gitlab': 'selected'}, db_factory=object, ) self.assertEqual(service.db_url, 'postgresql://fixture') self.assertIs(service.assignment_builder, builder.return_value) self.assertIsNone(service.admin_service) self.assertEqual(runtime, { 'address': '127.0.0.1', 'port': 8766, 'reaper_interval_seconds': 60, 'limit_concurrency': 64, }) def test_factory_defensively_rejects_invalid_assignment_ttl_overrides(self): for overrides in ( {'github': 3600}, {'gitlab': True}, {'gitlab': 59}, {'gitlab': 7 * 24 * 60 * 60 + 1}, ): with self.subTest(overrides=overrides): config = self._config( 'unused', assignment_ttl_seconds_by_source=overrides, ) with mock.patch.object( console_runner, 'load_config', return_value=config, ): with self.assertRaises(ValueError): worker_api.build_configured_worker_service( 'config.yaml', {'instance_id': 'fixture'}, ) with tempfile.TemporaryDirectory() as temp_dir: bundle_root = ensure_private_directory( os.path.join(temp_dir, 'bundles'), reject_reparse=True, ) config = self._config( bundle_root, assignment_ttl_seconds_by_source={'gitlab': 3659}, ) patches = self._patch_config(config) with patches[0], patches[1], patches[2], patches[3], patches[4], \ patches[5], patches[6]: with self.assertRaisesRegex(ValueError, 'for gitlab'): worker_api.build_configured_worker_service( 'config.yaml', {'instance_id': 'fixture'}, db_factory=object, ) def test_empty_source_list_builds_core_trio_without_direct_credentials(self): with tempfile.TemporaryDirectory() as temp_dir: bundle_root = ensure_private_directory( os.path.join(temp_dir, 'bundles'), reject_reparse=True, ) config = self._config(bundle_root, sources=[]) config['sources'] = { source: {'enabled': True} for source in ('gitlab', 'dockerhub', 'huggingface') } built = {} def build_args(source, *_args, auth_entry=None, **_kwargs): value = SimpleNamespace( platform={ 'gitlab': 'gitlab', 'dockerhub': 'docker', 'huggingface': 'huggingface', }[source], exact_git_planning_enabled=source == 'gitlab', timeout=30, token=(auth_entry or {}).get('token', 'ambient-token'), docker_username='ambient-user', docker_token='ambient-docker-token', auth_name='ambient-auth', ) built[source] = value return value with mock.patch.object(console_runner, 'load_config', return_value=config), \ mock.patch.object(console_runner, 'apply_global_config'), \ mock.patch.object(console_runner, 'load_secrets', return_value={'fixture': True}), \ mock.patch.object(console_runner, 'auth_pool_entries', return_value=([{ 'name': 'selected', 'token': 'selected-secret', }], None)) as auth_entries, \ mock.patch.object(console_runner, 'build_args_from_source_config', side_effect=build_args), \ mock.patch.object(db_backend, 'database_url_from_env', return_value='postgresql://fixture'), \ mock.patch.object(worker_assignment, 'RemoteAssignmentBuilder') as builder: worker_api.build_configured_worker_service( 'config.yaml', {'instance_id': 'supervisor-fixture'}, db_factory=object, ) auth_entries.assert_called_once() call = builder.call_args self.assertEqual(tuple(call.args[2]), ('gitlab', 'dockerhub', 'huggingface')) self.assertEqual(call.kwargs['credential_refs'], { 'gitlab': 'selected', 'dockerhub': '', 'huggingface': '', }) self.assertEqual(built['gitlab'].token, 'selected-secret') for source in ('dockerhub', 'huggingface'): self.assertEqual(built[source].token, '') self.assertEqual(built[source].docker_username, '') self.assertEqual(built[source].docker_token, '') self.assertIsNone(built[source].auth_name) def test_factory_rejects_direct_auth_entry_before_database_access(self): config = self._config( 'unused', sources=['dockerhub'], auth_entries={'dockerhub': 'private'}, ) with mock.patch.object(console_runner, 'load_config', return_value=config), \ mock.patch.object(db_backend, 'database_url_from_env') as database_url: with self.assertRaisesRegex(ValueError, 'only GitHub or GitLab'): worker_api.build_configured_worker_service( 'config.yaml', {'instance_id': 'fixture'}, ) database_url.assert_not_called() def test_legacy_github_auth_is_reconciliation_only(self): with tempfile.TemporaryDirectory() as temp_dir: bundle_root = ensure_private_directory( os.path.join(temp_dir, 'bundles'), reject_reparse=True, ) config = self._config( bundle_root, auth_entries={'gitlab': 'selected', 'github': 'legacy'}, ) config['sources']['github'] = {'enabled': True} built = {} def build_args(source, *_args, auth_entry=None, **_kwargs): value = SimpleNamespace( platform=source, exact_git_planning_enabled=True, timeout=30, token=(auth_entry or {}).get('token', ''), ) built[source] = value return value def auth_entries(_secrets, _source_config): return ([ {'name': 'selected', 'token': 'gitlab-secret'}, {'name': 'legacy', 'token': 'github-secret'}, ], None) with mock.patch.object(console_runner, 'load_config', return_value=config), \ mock.patch.object(console_runner, 'apply_global_config'), \ mock.patch.object(console_runner, 'load_secrets', return_value={'fixture': True}), \ mock.patch.object(console_runner, 'auth_pool_entries', side_effect=auth_entries), \ mock.patch.object(console_runner, 'build_args_from_source_config', side_effect=build_args), \ mock.patch.object(db_backend, 'database_url_from_env', return_value='postgresql://fixture'), \ mock.patch.object(worker_assignment, 'RemoteAssignmentBuilder') as builder: worker_api.build_configured_worker_service( 'config.yaml', {'instance_id': 'supervisor-fixture'}, db_factory=object, ) call = builder.call_args self.assertEqual(set(call.args[2]), {'gitlab', 'github'}) self.assertEqual(call.kwargs['credential_refs'], { 'gitlab': 'selected', 'github': 'legacy', }) self.assertEqual(built['github'].token, 'github-secret') self.assertEqual( config['supervisor']['worker_api']['sources'], ['gitlab'], ) def test_factory_builds_admin_only_from_enabled_nested_settings(self): with tempfile.TemporaryDirectory() as temp_dir: bundle_root = ensure_private_directory( os.path.join(temp_dir, 'bundles'), reject_reparse=True, ) config = self._config(bundle_root, admin={ 'enabled': True, 'origin': 'https://admin.example.test', 'edge_marker': 'fixture-edge-marker-value-32bytes-minimum', 'max_body_bytes': 4096, 'snapshot_limit': 17, 'requeue_limit': 9, 'managed_file_roots': { 'exports': { 'path': '/data/managed-files/exports', 'permissions': { 'list': True, 'read': True, 'create_replace': True, 'delete': True, }, 'limits': { 'max_relative_path_bytes': 1024, 'max_component_bytes': 255, 'max_path_depth': 16, 'max_listing_entries': 500, 'max_listing_bytes': 262144, 'max_file_bytes': 64 * 1024 * 1024, }, }, }, }) patches = self._patch_config(config) with patches[0], patches[1], patches[2], patches[3], patches[4], patches[5], \ patches[6] as builder, mock.patch.object( worker_api.HostAgentClient, 'is_available', return_value=True, ), mock.patch.object( worker_api, 'fixed_result_directory_is_safe', return_value=True, ): service, _runtime = worker_api.build_configured_worker_service( 'config.yaml', {'instance_id': 'supervisor-fixture'}, db_factory=object, ) self.assertEqual(service.admin_service.origin, 'https://admin.example.test') self.assertEqual(service.admin_service.max_body_bytes, 4096) self.assertEqual(service.admin_service.snapshot_limit, 17) self.assertEqual(service.admin_service.requeue_limit, 9) self.assertEqual( service.admin_service.managed_file_roots.root_ids(), ('exports',), ) self.assertEqual( service.admin_service.supervisor_metadata, {'instance_id': 'supervisor-fixture'}, ) self.assertIs( service.admin_service.package_compatibility_provider, builder.return_value.compatibility_snapshot, ) self.assertIsInstance( service.admin_service.runtime_apply_provider.__self__, worker_api.HostAgentClient, ) def test_factory_rejects_admin_unknown_keys(self): config = self._config('unused', admin={'enabled': False, 'command': 'restart'}) with mock.patch.object(console_runner, 'load_config', return_value=config): with self.assertRaisesRegex(ValueError, 'unsupported keys'): worker_api.build_configured_worker_service( 'config.yaml', {'instance_id': 'fixture'}, ) def test_factory_rejects_unsafe_managed_root_before_external_access(self): config = self._config('unused', admin={ 'enabled': False, 'managed_file_roots': { 'private': { 'path': '/data/config', 'permissions': { 'list': True, 'read': True, 'create_replace': False, 'delete': False, }, 'limits': { 'max_relative_path_bytes': 1024, 'max_component_bytes': 255, 'max_path_depth': 16, 'max_listing_entries': 500, 'max_listing_bytes': 262144, 'max_file_bytes': 64 * 1024 * 1024, }, }, }, }) with mock.patch.object(console_runner, 'load_config', return_value=config), \ mock.patch.object(console_runner, 'load_secrets') as load_secrets, \ mock.patch.object(db_backend, 'database_url_from_env') as database_url, \ mock.patch.object(worker_assignment, 'RemoteAssignmentBuilder') as builder: with self.assertRaises(managed_files.ManagedFileConfigurationError): worker_api.build_configured_worker_service( 'config.yaml', {'instance_id': 'fixture'}, ) load_secrets.assert_not_called() database_url.assert_not_called() builder.assert_not_called() def test_main_authenticates_before_building_and_never_enables_proxy_headers(self): trace = [] metadata = {'instance_id': 'fixture'} runtime = { 'address': '127.0.0.1', 'port': 8766, 'reaper_interval_seconds': 60, 'limit_concurrency': 64, } def authenticate(*args, **kwargs): trace.append('authenticate') return metadata def build(*args, **kwargs): trace.append('build') return object(), runtime with mock.patch.object(worker_api, 'parse_args', return_value=SimpleNamespace(config='config.yaml')), \ mock.patch.object(lifecycle_authority, 'require_active_supervisor_child', side_effect=authenticate) as authority, \ mock.patch.object(worker_api, 'build_configured_worker_service', side_effect=build), \ mock.patch.object(worker_api, 'create_worker_app', return_value='fixture-app'), \ mock.patch.object(uvicorn, 'run') as run: worker_api.main() self.assertEqual(trace, ['authenticate', 'build']) authority.assert_called_once_with( 'config.yaml', child_kind='worker-api', require_dsn=True, ) run.assert_called_once_with( 'fixture-app', host='127.0.0.1', port=8766, access_log=False, proxy_headers=False, server_header=False, limit_concurrency=64, timeout_keep_alive=5, workers=1, ) class WorkerSupervisorWiringTests(unittest.TestCase): def test_child_bootstrap_dispatches_worker_api_from_immutable_authority(self): namespace = runpy.run_path(str(APP_DIR / 'child_bootstrap.py')) main = namespace['main'] entrypoint = os.path.normcase(os.path.realpath(APP_DIR / 'worker_api.py')) metadata = {'code_manifest': {'files': {'worker-api': {'path': entrypoint}}}} enabled = [] captured = [] with mock.patch.object(sys, 'flags', SimpleNamespace( isolated=1, no_site=1, dont_write_bytecode=1, )), mock.patch.object(sys, 'path', list(sys.path)), mock.patch.object( sys, 'argv', [str(APP_DIR / 'child_bootstrap.py'), 'worker-api', '--', '--config', 'config.yaml'], ), mock.patch.dict(main.__globals__, { '_authenticate': lambda kind: (str(APP_DIR), metadata), '_enable_dependency_paths': lambda kind: enabled.append(kind), }), mock.patch.object( main.__globals__['runpy'], 'run_path', side_effect=lambda path, run_name=None: captured.append((path, list(sys.argv), run_name)), ): main() self.assertEqual(enabled, ['worker-api']) self.assertEqual(captured, [( entrypoint, [entrypoint, '--config', 'config.yaml'], '__main__', )]) def _worker(self, enabled): dependency = mock.Mock() with mock.patch.object(supervisor, 'require_private_directory', side_effect=lambda path, create=False: path): worker = supervisor.ManagedPipelineWorker( 'worker-api', 'config.yaml', 'project', 'results', {}, {'enabled': enabled} if enabled is not None else {}, dependency_gate=dependency, child_environment={'TRUF_SUPERVISOR_CHILD_KIND': 'worker-api'}, ) return worker, dependency def test_worker_api_child_defaults_disabled_and_uses_authenticated_bootstrap(self): disabled, _ = self._worker(None) self.assertFalse(disabled.enabled) self.assertEqual(disabled.status, 'disabled') enabled, dependency = self._worker(True) self.assertTrue(enabled.enabled) command = enabled.build_command() self.assertEqual(command[:4], [sys.executable, '-I', '-S', '-B']) self.assertIn('worker-api', command) self.assertEqual(command[-2:], ['--config', enabled.config_path]) environment = enabled.build_env() dependency.force_database_environment.assert_called_once_with(environment) self.assertEqual(environment['TRUF_SUPERVISOR_CHILD_KIND'], 'worker-api') if __name__ == '__main__': unittest.main()