Files
truf-server/tests/test_worker_observability_db.py
2026-09-30 20:30:56 +03:00

644 lines
28 KiB
Python

import hashlib
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 scanner_db
import worker_contracts
from scanner_db import (
ScannerDB,
WorkerObservabilityConflictError,
WorkerProgressInactiveError,
)
def canonical_contract(_kind, payload):
normalized = dict(payload)
encoded = json.dumps(
normalized, ensure_ascii=True, sort_keys=True, separators=(',', ':'),
)
return normalized, encoded, hashlib.sha256(encoded.encode('utf-8')).hexdigest()
class WorkerObservabilityDbTests(unittest.TestCase):
def setUp(self):
self.environment = mock.patch.dict(
os.environ, {'SCANNER_DB_URL': '', 'DATABASE_URL': ''},
)
self.environment.start()
self.temp = tempfile.TemporaryDirectory()
self.db = ScannerDB(db_path=os.path.join(self.temp.name, 'scanner.db'))
self.contract = mock.patch.object(
scanner_db, '_canonical_worker_contract', side_effect=canonical_contract,
)
self.contract.start()
(
self.device_id, self.queue_id, self.reservation_id, self.scan_event_id,
) = self._seed_reservation('one')
def tearDown(self):
self.contract.stop()
self.db.close()
self.temp.cleanup()
self.environment.stop()
def _seed_reservation(self, label):
now = '2026-09-23T12:00:00+00:00'
scan_event_id = hashlib.sha256(f'scan:{label}'.encode('ascii')).hexdigest()[:32]
user_id = self.db.conn.insert_returning_id(
'''INSERT INTO remote_worker_users(
user_key, active_assignment_cap, created_at, updated_at
) VALUES (?, 1, ?, ?)''',
(f'user-{label}', now, now),
)
device_id = self.db.conn.insert_returning_id(
'''INSERT INTO remote_worker_devices(
user_id, device_key, token_sha256, created_at, updated_at
) VALUES (?, ?, ?, ?, ?)''',
(user_id, f'device-{label}', hashlib.sha256(label.encode()).hexdigest(), now, now),
)
target = f'https://github.com/example/{label}.git'
queue_id = self.db.conn.insert_returning_id(
'''INSERT INTO target_queue(
source, platform, query, target, normalized_target, status,
lease_owner, lease_token, claim_batch, leased_at, lease_expires_at,
claim_event_id, created_at, updated_at
) VALUES ('github', 'github', 'fixture', ?, ?, 'in_progress',
'worker-api', ?, ?, ?, '2999-01-01T00:00:00+00:00',
?, ?, ?)''',
(target, target, f'lease-{label}', f'batch-{label}', now,
scan_event_id, now, now),
)
reservation_id = self.db.conn.insert_returning_id(
'''INSERT INTO result_reservations(
reservation_token, bundle_id, scan_event_id, queue_id,
source, platform, query, target, normalized_target,
claim_lease_owner, claim_lease_token, claim_batch,
producer_instance_id, producer_pid, producer_creation_time,
producer_executable, assignment_kind, remote_user_id,
remote_device_id, remote_issued_at, remote_expires_at,
declared_bundle_bytes, reserved_bundle_bytes,
reserved_projection_bytes,
reserved_candidate_items, reserved_candidate_bytes,
ready_relative_path, state, producer_lease_expires_at,
created_at, updated_at
) VALUES (?, ?, ?, ?, 'github', 'github', 'fixture', ?, ?,
'worker-api', ?, ?, 'api-instance', 1, ?, 'worker-api',
'remote', ?, ?, ?, '2999-01-01T00:00:00+00:00',
1024, 1024, 1024, 0, 0, ?, 'scanning',
'2999-01-01T00:00:00+00:00', ?, ?)''',
(
f'reservation-{label}', f'bundle-{label}', scan_event_id, queue_id,
target, target, f'lease-{label}', f'batch-{label}', now,
user_id, device_id, now, f'ready/{label}.bundle', now, now,
),
)
self.db.conn.execute(
'''UPDATE target_queue SET current_result_reservation_id = ?
WHERE id = ?''',
(reservation_id, queue_id),
)
self.db.conn.commit()
return device_id, queue_id, reservation_id, scan_event_id
def _event(self, sequence, phase='scanning'):
return {
'schema': 1,
'sequence': sequence,
'timestamp': f'2026-09-23T12:00:{sequence % 60:02d}Z',
'instance_id': 'worker-instance',
'slot_id': 0,
'reservation_id': self.reservation_id,
'source': 'github',
'type': 'slot.phase',
'phase': phase,
'phase_started_at': '2026-09-23T12:00:00Z',
'progress': {},
}
def _diagnostic(self, uid='diag-one', summary='scanner failed'):
return {
'schema': 1,
'diagnostic_uid': uid,
'reservation_id': self.reservation_id,
'scan_event_id': self.scan_event_id,
'slot_id': 0,
'attempt': 1,
'source': 'github',
'phase': 'scanning',
'kind': 'process',
'category': 'scanner',
'code': 'trufflehog.exit_nonzero',
'summary': summary,
'retryable': False,
'assignment_outcome': 'accepted',
'scan_outcome': 'error',
'occurred_at': '2026-09-23T12:01:00Z',
'captured_at': '2026-09-23T12:01:01Z',
'body': {'encoding': 'utf-8', 'text': 'provider body', 'truncated': False},
'logs': {'stderr': {'text': 'failure', 'truncated': False}},
}
def _target_scan(self):
scan_id = self.db.conn.insert_returning_id(
'''INSERT INTO target_scans(
scan_event_id, result_reservation_id, source, target, status,
started_at, ended_at, created_at
) VALUES (?, ?, 'github', 'fixture', 'error', ?, ?, ?)''',
(
self.scan_event_id, self.reservation_id, '2026-09-23T12:00:00Z',
'2026-09-23T12:01:00Z', '2026-09-23T12:01:00Z',
),
)
self.db.conn.commit()
return scan_id
def test_schema_migration_indexes_and_foreign_keys_are_authoritative(self):
self.assertEqual(
scanner_db.PIPELINE_MIGRATION_VERSIONS[-1],
scanner_db.REMOTE_ASSIGNMENT_CAPACITY_MIGRATION,
)
self.assertEqual(len(scanner_db.PIPELINE_MIGRATION_VERSIONS), 33)
self.assertIn(
'remote_result_upload_body_timeout_seconds',
self.db.conn.table_columns('result_reservations'),
)
self.assertIn(
'reserved_bundle_bytes',
self.db.conn.table_columns('result_reservations'),
)
for table, specs in (
('worker_progress_events', scanner_db.WORKER_PROGRESS_EVENT_COLUMN_SPECS),
('worker_diagnostics', scanner_db.WORKER_DIAGNOSTIC_COLUMN_SPECS),
):
with self.subTest(table=table):
self.assertEqual(set(self.db.conn.table_columns(table)), set(specs))
foreign_keys = self.db.conn.table_foreign_keys(table)
for columns, referenced_table, referenced_columns in scanner_db.REQUIRED_FOREIGN_KEYS[table]:
self.assertTrue(any(
tuple(item['columns']) == columns
and item['referenced_table'] == referenced_table
and tuple(item['referenced_columns']) == referenced_columns
for item in foreign_keys.values()
))
progress_indexes = self.db.conn.table_indexes('worker_progress_events')
diagnostic_indexes = self.db.conn.table_indexes('worker_diagnostics')
self.assertTrue(progress_indexes['uq_worker_progress_reservation_sequence']['unique'])
self.assertTrue(diagnostic_indexes['uq_worker_diagnostics_uid']['unique'])
self.assertTrue(self.db.runtime_safety_schema_available(), self.db.last_error)
self.assertTrue(self.db.pipeline_schema_available(), self.db.last_error)
def test_additive_migration_recreates_observability_authority(self):
with tempfile.TemporaryDirectory() as temp:
db = ScannerDB(db_path=os.path.join(temp, 'pre-observability.db'))
try:
db.conn.execute('DROP TABLE worker_diagnostics')
db.conn.execute('DROP TABLE worker_progress_events')
db.conn.execute('DROP INDEX idx_target_scans_result_reservation')
db.conn.execute(
'DELETE FROM runtime_schema_migrations WHERE version = ?',
(scanner_db.WORKER_OBSERVABILITY_MIGRATION,),
)
db.conn.commit()
scanner_db.migrate_runtime_safety_schema(db)
self.assertTrue(db.conn.table_exists('worker_progress_events'))
self.assertTrue(db.conn.table_exists('worker_diagnostics'))
self.assertIn(
'idx_target_scans_result_reservation',
db.conn.table_indexes('target_scans'),
)
marker = db.conn.execute(
'SELECT code_sha256 FROM runtime_schema_migrations WHERE version = ?',
(scanner_db.WORKER_OBSERVABILITY_MIGRATION,),
).fetchone()
self.assertIsNotNone(marker)
self.assertTrue(db.runtime_safety_schema_available(), db.last_error)
finally:
db.close()
def test_remote_capacity_migration_reconciles_active_persisted_charges(self):
self.db.conn.execute(
'''UPDATE pipeline_capacity SET bundle_items = 99, bundle_bytes = 99,
projection_items = 99, projection_bytes = 99,
keycheck_items = 99, keycheck_bytes = 99 WHERE id = 1'''
)
self.db.conn.execute(
'DELETE FROM runtime_schema_migrations WHERE version = ?',
(scanner_db.REMOTE_ASSIGNMENT_CAPACITY_MIGRATION,),
)
self.db.conn.commit()
scanner_db.migrate_runtime_safety_schema(self.db)
capacity = self.db.pipeline_capacity_snapshot()
self.assertEqual(
(
capacity['bundle_items'], capacity['bundle_bytes'],
capacity['projection_items'], capacity['projection_bytes'],
capacity['keycheck_items'], capacity['keycheck_bytes'],
),
(1, 1024, 1, 1024, 0, 0),
)
self.assertIsNotNone(self.db.conn.execute(
'SELECT version FROM runtime_schema_migrations WHERE version = ?',
(scanner_db.REMOTE_ASSIGNMENT_CAPACITY_MIGRATION,),
).fetchone())
def test_projection_authority_migration_backfills_old_acceptance_as_unfenced(self):
self.db.conn.execute(
'''UPDATE result_reservations
SET state = 'acknowledged',
remote_resolution_kind = 'bundle_accepted',
remote_diagnostic_projection_version = NULL,
remote_diagnostic_count = NULL,
remote_diagnostic_uids_sha256 = NULL
WHERE id = ?''',
(self.reservation_id,),
)
self.db.conn.execute(
'''UPDATE target_queue SET status = 'done', lease_token = NULL,
current_result_reservation_id = NULL, claim_event_id = NULL
WHERE id = ?''',
(self.queue_id,),
)
self.db.conn.execute(
'DELETE FROM runtime_schema_migrations WHERE version = ?',
(scanner_db.DIAGNOSTIC_PROJECTION_AUTHORITY_MIGRATION,),
)
self.db.conn.commit()
scanner_db.migrate_runtime_safety_schema(self.db)
row = self.db.conn.execute(
'''SELECT remote_diagnostic_projection_version,
remote_diagnostic_count, remote_diagnostic_uids_sha256
FROM result_reservations WHERE id = ?''',
(self.reservation_id,),
).fetchone()
self.assertEqual(row['remote_diagnostic_projection_version'], 0)
self.assertIsNone(row['remote_diagnostic_count'])
self.assertIsNone(row['remote_diagnostic_uids_sha256'])
marker = self.db.conn.execute(
'SELECT version FROM runtime_schema_migrations WHERE version = ?',
(scanner_db.DIAGNOSTIC_PROJECTION_AUTHORITY_MIGRATION,),
).fetchone()
self.assertIsNotNone(marker)
def test_upload_timeout_migration_preserves_legacy_rows_as_unavailable(self):
self.db.conn.execute(
'''UPDATE result_reservations
SET remote_result_upload_body_timeout_seconds = NULL,
state = 'acknowledged', remote_resolution_kind = 'expired'
WHERE id = ?''',
(self.reservation_id,),
)
self.db.conn.execute(
'''UPDATE target_queue SET status = 'done', lease_token = NULL,
current_result_reservation_id = NULL, claim_event_id = NULL
WHERE id = ?''',
(self.queue_id,),
)
self.db.conn.execute(
'DELETE FROM runtime_schema_migrations WHERE version = ?',
(scanner_db.ASSIGNMENT_UPLOAD_TIMEOUT_MIGRATION,),
)
self.db.conn.commit()
scanner_db.migrate_runtime_safety_schema(self.db)
row = self.db.conn.execute(
'''SELECT remote_result_upload_body_timeout_seconds
FROM result_reservations WHERE id = ?''',
(self.reservation_id,),
).fetchone()
self.assertIsNone(row['remote_result_upload_body_timeout_seconds'])
marker = self.db.conn.execute(
'SELECT version FROM runtime_schema_migrations WHERE version = ?',
(scanner_db.ASSIGNMENT_UPLOAD_TIMEOUT_MIGRATION,),
).fetchone()
self.assertIsNotNone(marker)
reservation = self.db.conn.execute(
'SELECT * FROM result_reservations WHERE id = ?',
(self.reservation_id,),
).fetchone()
observability = self.db._remote_assignment_observability_locked(
reservation,
)
self.assertIsNone(
observability['deadlines']['result_upload_body_timeout_seconds'],
)
self.assertEqual(
observability['deadlines'][
'result_upload_body_timeout_availability'
],
'legacy/unavailable',
)
def test_progress_is_strict_idempotent_and_never_renews_authority(self):
before = dict(self.db.conn.execute(
'''SELECT r.remote_expires_at, r.producer_lease_expires_at,
q.lease_expires_at
FROM result_reservations r JOIN target_queue q ON q.id = r.queue_id
WHERE r.id = ?''',
(self.reservation_id,),
).fetchone())
first = self.db._record_worker_progress_event(
self.reservation_id, self.device_id, self._event(41),
)
self.assertFalse(first['replayed'])
with self.assertRaisesRegex(
WorkerObservabilityConflictError, 'owned reservation',
):
self.db._record_worker_progress_event(
self.reservation_id, self.device_id + 1, self._event(42),
)
replay = self.db._record_worker_progress_event(
self.reservation_id, self.device_id, self._event(41),
)
self.assertTrue(replay['replayed'])
with self.assertRaisesRegex(WorkerObservabilityConflictError, 'different canonical'):
self.db._record_worker_progress_event(
self.reservation_id, self.device_id, self._event(41, phase='bundling'),
)
self.db._record_worker_progress_event(
self.reservation_id, self.device_id, self._event(43, phase='bundling'),
)
with self.assertRaisesRegex(WorkerObservabilityConflictError, 'strictly increase'):
self.db._record_worker_progress_event(
self.reservation_id, self.device_id, self._event(42),
)
after = dict(self.db.conn.execute(
'''SELECT r.remote_expires_at, r.producer_lease_expires_at,
q.lease_expires_at
FROM result_reservations r JOIN target_queue q ON q.id = r.queue_id
WHERE r.id = ?''',
(self.reservation_id,),
).fetchone())
self.assertEqual(after, before)
self.db.conn.execute(
'UPDATE target_queue SET current_result_reservation_id = NULL WHERE id = ?',
(self.queue_id,),
)
replay_after_settlement = self.db._record_worker_progress_event(
self.reservation_id, self.device_id, self._event(41),
)
self.assertTrue(replay_after_settlement['replayed'])
with self.assertRaisesRegex(WorkerObservabilityConflictError, 'owned current unresolved'):
self.db._record_worker_progress_event(
self.reservation_id, self.device_id, self._event(44),
)
receipt = {
'receipt_id': 'a' * 64,
'resolution': 'prebundle_report',
}
receipt_json = json.dumps(
receipt, ensure_ascii=True, sort_keys=True, separators=(',', ':'),
)
self.db.conn.execute(
'''UPDATE result_reservations SET state = 'refunded',
remote_resolution_kind = 'prebundle_report',
remote_receipt_id = ?, remote_resolution_json = ?,
remote_resolved_at = ? WHERE id = ?''',
('a' * 64, receipt_json, '2026-09-23T12:02:00+00:00', self.reservation_id),
)
authority_before = dict(self.db.conn.execute(
'''SELECT remote_expires_at, producer_lease_expires_at,
remote_receipt_id, remote_resolution_json
FROM result_reservations WHERE id = ?''',
(self.reservation_id,),
).fetchone())
late = self.db._record_worker_progress_event(
self.reservation_id, self.device_id,
self._event(45, phase='awaiting_receipt'),
received_at='2026-09-23T12:02:30+00:00',
)
replayed_late = self.db._record_worker_progress_event(
self.reservation_id, self.device_id,
self._event(45, phase='awaiting_receipt'),
received_at='2026-09-24T12:02:30+00:00',
)
authority_after = dict(self.db.conn.execute(
'''SELECT remote_expires_at, producer_lease_expires_at,
remote_receipt_id, remote_resolution_json
FROM result_reservations WHERE id = ?''',
(self.reservation_id,),
).fetchone())
self.assertFalse(late['replayed'])
self.assertTrue(replayed_late['replayed'])
self.assertEqual(authority_after, authority_before)
def test_diagnostic_helper_is_commit_free_and_replay_safe(self):
scan_id = self._target_scan()
inserted = self.db._record_worker_diagnostic(
self.reservation_id, self._diagnostic(), target_scan_id=scan_id,
)
self.assertFalse(inserted['replayed'])
self.assertEqual(json.loads(inserted['body_payload_json'])['text'], 'provider body')
self.db.conn.rollback()
count = self.db.conn.execute(
'SELECT COUNT(*) AS count FROM worker_diagnostics'
).fetchone()
self.assertEqual(int(count['count']), 0)
inserted = self.db._record_worker_diagnostic(
self.reservation_id, self._diagnostic(), target_scan_id=scan_id,
)
replay = self.db._record_worker_diagnostic(
self.reservation_id, self._diagnostic(), target_scan_id=scan_id,
)
self.assertFalse(inserted['replayed'])
self.assertTrue(replay['replayed'])
with self.assertRaisesRegex(WorkerObservabilityConflictError, 'different canonical'):
self.db._record_worker_diagnostic(
self.reservation_id,
self._diagnostic(summary='different failure'),
target_scan_id=scan_id,
)
self.db.conn.commit()
rows = self.db.worker_diagnostics(self.reservation_id)
self.assertEqual(len(rows), 1)
self.assertEqual(rows[0]['target_scan_id'], scan_id)
self.assertEqual(rows[0]['scan_event_id'], self.scan_event_id)
self.assertEqual(rows[0]['scan_status'], 'error')
self.assertEqual(rows[0]['queue_status'], 'in_progress')
def test_progress_query_joins_reservation_queue_and_scan(self):
scan_id = self._target_scan()
self.db._record_worker_progress_event(
self.reservation_id, self.device_id, self._event(7),
)
self.db.conn.commit()
rows = self.db.worker_progress_events(self.reservation_id)
self.assertEqual(len(rows), 1)
self.assertEqual(rows[0]['target_scan_id'], scan_id)
self.assertEqual(rows[0]['scan_status'], 'error')
self.assertEqual(rows[0]['queue_id'], self.queue_id)
self.assertEqual(rows[0]['event']['phase'], 'scanning')
def test_resolved_progress_grace_and_timestamp_windows_are_bounded(self):
with self.assertRaisesRegex(
WorkerObservabilityConflictError, 'authority window',
):
self.db._record_worker_progress_event(
self.reservation_id, self.device_id,
self._event(1) | {'timestamp': '2026-09-23T11:00:00Z'},
received_at='2026-09-23T12:00:01+00:00',
)
with self.assertRaisesRegex(
WorkerObservabilityConflictError, 'authority window',
):
self.db._record_worker_progress_event(
self.reservation_id, self.device_id,
self._event(2) | {'timestamp': '2026-09-23T13:00:00Z'},
received_at='2026-09-23T12:00:01+00:00',
)
receipt = {'receipt_id': 'b' * 64, 'resolution': 'prebundle_report'}
self.db.conn.execute(
'''UPDATE result_reservations SET state = 'refunded',
remote_resolution_kind = 'prebundle_report',
remote_receipt_id = ?, remote_resolution_json = ?,
remote_resolved_at = ? WHERE id = ?''',
(
'b' * 64,
json.dumps(receipt, sort_keys=True, separators=(',', ':')),
'2026-09-23T12:05:00+00:00', self.reservation_id,
),
)
self.db.conn.execute(
'''UPDATE target_queue SET status = 'pending', lease_token = NULL,
current_result_reservation_id = NULL, claim_event_id = NULL
WHERE id = ?''',
(self.queue_id,),
)
with self.assertRaisesRegex(
WorkerProgressInactiveError, 'too old for terminal grace',
):
self.db._record_worker_progress_event(
self.reservation_id, self.device_id,
self._event(9, phase='uploading') | {
'timestamp': '2026-09-23T12:01:59Z',
'phase_started_at': '2026-09-23T12:01:59Z',
},
received_at='2026-09-23T12:05:10+00:00',
)
accepted = self.db._record_worker_progress_event(
self.reservation_id, self.device_id,
self._event(10, phase='awaiting_receipt') | {
'timestamp': '2026-09-23T12:04:59Z',
'phase_started_at': '2026-09-23T12:04:59Z',
},
received_at='2026-09-23T12:06:59+00:00',
)
replay = self.db._record_worker_progress_event(
self.reservation_id, self.device_id,
accepted['event'],
received_at='2026-09-24T12:06:59+00:00',
)
self.assertTrue(replay['replayed'])
with self.assertRaisesRegex(
WorkerProgressInactiveError, 'grace window',
):
self.db._record_worker_progress_event(
self.reservation_id, self.device_id,
self._event(11, phase='awaiting_receipt') | {
'timestamp': '2026-09-23T12:05:00Z',
'phase_started_at': '2026-09-23T12:05:00Z',
},
received_at='2026-09-23T12:07:01+00:00',
)
accepted_skew = self.db._record_worker_progress_event(
self.reservation_id, self.device_id,
self._event(12, phase='awaiting_receipt') | {
'timestamp': '2026-09-23T12:05:22Z',
'phase_started_at': '2026-09-23T12:05:22Z',
},
received_at='2026-09-23T12:05:30+00:00',
)
self.assertEqual(accepted_skew['sequence'], 12)
with self.assertRaisesRegex(
WorkerObservabilityConflictError, 'after terminal resolution',
):
self.db._record_worker_progress_event(
self.reservation_id, self.device_id,
self._event(13, phase='awaiting_receipt') | {
'timestamp': '2026-09-23T12:06:01Z',
'phase_started_at': '2026-09-23T12:06:01Z',
},
received_at='2026-09-23T12:06:10+00:00',
)
def test_missing_contract_module_fails_with_actionable_error(self):
self.contract.stop()
try:
missing = ModuleNotFoundError(name='worker_contracts')
with mock.patch.object(scanner_db.importlib, 'import_module', side_effect=missing):
with self.assertRaisesRegex(
RuntimeError,
'worker_contracts is required to persist worker progress records',
):
scanner_db._canonical_worker_contract('progress', self._event(1))
finally:
self.contract.start()
def test_available_contract_codecs_are_used_for_canonical_validation(self):
self.contract.stop()
try:
event = self._event(1)
event.update(scan_deadline_at=None, assignment_deadline_at=None)
normalized, canonical, digest = scanner_db._canonical_worker_contract(
'progress', event,
)
self.assertEqual(normalized, event)
self.assertEqual(
digest, hashlib.sha256(canonical.encode('ascii')).hexdigest(),
)
timestamp = '2026-09-23T12:00:00Z'
envelope = worker_contracts.build_diagnostic_envelope(
occurrence_id='contract-codec',
reservation_id=self.reservation_id,
scan_event_id=None,
slot_id=0,
source='github',
phase=worker_contracts.WorkerPhase.SCANNING,
kind=worker_contracts.DiagnosticKind.ASSIGNMENT,
category=worker_contracts.DiagnosticCategory.INTERNAL,
code='worker.contract_codec',
summary='contract codec fixture',
retryable=False,
attempt=1,
assignment_outcome=None,
scan_outcome=None,
occurred_at=timestamp,
captured_at=timestamp,
)
diagnostic = json.loads(
worker_contracts.encode_diagnostic_envelope(envelope).decode('ascii')
)
normalized, canonical, digest = scanner_db._canonical_worker_contract(
'diagnostic', diagnostic,
)
self.assertEqual(normalized, diagnostic)
self.assertEqual(
digest, hashlib.sha256(canonical.encode('ascii')).hexdigest(),
)
finally:
self.contract.start()
if __name__ == '__main__':
unittest.main()