sashatrask
477 lines
22 KiB
Python
477 lines
22 KiB
Python
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()
|