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