238 lines
12 KiB
Python
238 lines
12 KiB
Python
import argparse
|
|
from contextlib import redirect_stderr, redirect_stdout
|
|
import fcntl
|
|
import io
|
|
import json
|
|
import os
|
|
import signal
|
|
from pathlib import Path
|
|
import tempfile
|
|
import unittest
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
from car_tracker import pipeline
|
|
|
|
|
|
class PipelineTests(unittest.TestCase):
|
|
def setUp(self):
|
|
self.directory = tempfile.TemporaryDirectory()
|
|
self.addCleanup(self.directory.cleanup)
|
|
self.output = Path(self.directory.name) / "raw"
|
|
self.args = argparse.Namespace(output=self.output, searches=[])
|
|
self.stdout, self.stderr = io.StringIO(), io.StringIO()
|
|
self.enterContext(redirect_stdout(self.stdout))
|
|
self.enterContext(redirect_stderr(self.stderr))
|
|
self.enterContext(patch.dict(os.environ, {"PGHOST": "homelab"}, clear=True))
|
|
self.preflight, self.import_connection = MagicMock(), MagicMock()
|
|
self.preflight.__enter__.return_value = self.preflight
|
|
self.import_connection.__enter__.return_value = self.import_connection
|
|
self.connect = self.enterContext(patch.object(
|
|
pipeline.psycopg, "connect", side_effect=[self.preflight, self.import_connection],
|
|
))
|
|
self.collect = self.enterContext(patch.object(pipeline, "collect_browser", side_effect=self.fixture))
|
|
self.summary = {"status": "imported", "accepted": 1, "rejected": 0}
|
|
self.importer = self.enterContext(patch.object(pipeline, "import_run", return_value=self.summary))
|
|
self.path = self.output / "returned-run" / "manifest.json"
|
|
self.raw = self.path.parent / "search-001.jsonl.gz"
|
|
|
|
def fixture(self, args, status="complete", error_type=None):
|
|
self.preflight.__exit__.assert_called_once()
|
|
self.path.parent.mkdir(parents=True, exist_ok=True)
|
|
manifest = {"status": status}
|
|
if error_type:
|
|
manifest["error_type"] = error_type
|
|
self.path.write_text(json.dumps(manifest), encoding="utf-8")
|
|
self.raw.write_bytes(b"raw snapshot")
|
|
self.manifest_bytes = self.path.read_bytes()
|
|
return self.path, status
|
|
|
|
def state(self):
|
|
return json.loads((self.path.parent / "pipeline.json").read_text())
|
|
|
|
def test_collection_statuses_import_exact_manifest(self):
|
|
for status, code in (("complete", 0), ("partial", 1), ("failed", 1)):
|
|
with self.subTest(status=status):
|
|
self.preflight.reset_mock()
|
|
self.connect.side_effect = [self.preflight, self.import_connection]
|
|
self.collect.side_effect = lambda args: self.fixture(args, status)
|
|
stale = self.output / "zz-newer-run" / "manifest.json"
|
|
stale.parent.mkdir(parents=True, exist_ok=True)
|
|
stale.write_text('{"status": "complete"}')
|
|
output_mode = self.output.stat().st_mode
|
|
self.assertEqual(pipeline.run_pipeline(self.args), code)
|
|
self.assertEqual(self.output.stat().st_mode, output_mode)
|
|
self.importer.assert_called_with(self.import_connection, self.path)
|
|
self.assertEqual(self.path.read_bytes(), self.manifest_bytes)
|
|
self.assertEqual(self.state()["collection_status"], status)
|
|
self.assertEqual(self.state()["import_status"], "succeeded")
|
|
self.preflight.execute.assert_called_once_with("SELECT 1 FROM silver.collection_runs LIMIT 0")
|
|
for call in self.connect.call_args_list:
|
|
self.assertEqual(call.args, ("",))
|
|
self.assertEqual(call.kwargs, {"connect_timeout": 10, "autocommit": True})
|
|
|
|
def test_new_output_directory_is_private(self):
|
|
self.assertEqual(pipeline.run_pipeline(self.args), 0)
|
|
self.assertEqual(self.output.stat().st_mode & 0o777, 0o700)
|
|
|
|
def test_rejections_take_priority_over_partial_collection(self):
|
|
for status in ("complete", "partial", "failed"):
|
|
with self.subTest(status=status):
|
|
self.preflight.reset_mock()
|
|
self.connect.side_effect = [self.preflight, self.import_connection]
|
|
self.collect.side_effect = lambda args: self.fixture(args, status)
|
|
self.summary["rejected"] = 1
|
|
self.assertEqual(pipeline.run_pipeline(self.args), 2)
|
|
|
|
def test_import_failure_preserves_raw_and_reports_promptly_without_secrets(self):
|
|
def fail(conn, path):
|
|
self.assertIn(str(path), self.stdout.getvalue())
|
|
raise pipeline.psycopg.OperationalError("password authentication failed TOP_SECRET")
|
|
self.importer.side_effect = fail
|
|
self.assertEqual(pipeline.run_pipeline(self.args), 3)
|
|
self.assertEqual(self.raw.read_bytes(), b"raw snapshot")
|
|
self.assertEqual(self.path.read_bytes(), self.manifest_bytes)
|
|
self.assertEqual(self.state()["import_status"], "unknown")
|
|
self.assertEqual(self.state()["error_type"], "OperationalError")
|
|
self.assertNotIn("TOP_SECRET", self.stderr.getvalue() + self.stdout.getvalue() + json.dumps(self.state()))
|
|
|
|
def test_preflight_failure_never_collects(self):
|
|
self.preflight.execute.side_effect = pipeline.psycopg.OperationalError("TOP_SECRET")
|
|
self.assertEqual(pipeline.run_pipeline(self.args), 3)
|
|
self.collect.assert_not_called()
|
|
self.importer.assert_not_called()
|
|
self.preflight.__exit__.assert_called_once()
|
|
self.assertNotIn("TOP_SECRET", self.stderr.getvalue())
|
|
|
|
def test_incremental_history_is_loaded_before_collection(self):
|
|
self.args.detail_policy = "incremental"
|
|
self.args.detail_refresh_days = 7
|
|
with patch.object(pipeline, "load_policy", return_value="synthetic-policy") as load:
|
|
self.assertEqual(pipeline.run_pipeline(self.args), 0)
|
|
load.assert_called_once_with(self.preflight, 7)
|
|
self.assertEqual(self.args.detail_decider, "synthetic-policy")
|
|
|
|
def test_history_failure_does_not_start_collection(self):
|
|
self.args.detail_policy = "incremental"
|
|
self.args.detail_refresh_days = 7
|
|
with patch.object(pipeline, "load_policy", side_effect=pipeline.psycopg.OperationalError("secret")):
|
|
self.assertEqual(pipeline.run_pipeline(self.args), 3)
|
|
self.collect.assert_not_called()
|
|
|
|
def test_finalized_interruption_is_not_imported(self):
|
|
self.collect.side_effect = lambda args: self.fixture(args, "partial", "KeyboardInterrupt")
|
|
self.assertEqual(pipeline.run_pipeline(self.args), 130)
|
|
self.importer.assert_not_called()
|
|
self.assertEqual(self.state()["import_status"], "not_attempted")
|
|
self.assertEqual(self.state()["error_type"], "KeyboardInterrupt")
|
|
self.assertEqual(self.path.read_bytes(), self.manifest_bytes)
|
|
|
|
def test_interruption_before_collector_returns(self):
|
|
self.collect.side_effect = KeyboardInterrupt
|
|
self.assertEqual(pipeline.run_pipeline(self.args), 130)
|
|
self.importer.assert_not_called()
|
|
|
|
def test_interruption_during_import_exits_connection_with_interrupt(self):
|
|
self.importer.side_effect = KeyboardInterrupt
|
|
self.assertEqual(pipeline.run_pipeline(self.args), 130)
|
|
self.assertIs(self.import_connection.__exit__.call_args.args[0], KeyboardInterrupt)
|
|
self.assertEqual(self.state()["import_status"], "unknown")
|
|
|
|
def test_lock_refusal_and_release_keeps_inode(self):
|
|
self.output.mkdir()
|
|
lock_path = self.output / ".pipeline.lock"
|
|
with lock_path.open("a") as lock:
|
|
inode = lock_path.stat().st_ino
|
|
fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB)
|
|
self.assertEqual(pipeline.run_pipeline(self.args), 3)
|
|
self.collect.assert_not_called()
|
|
self.connect.assert_not_called()
|
|
self.collect.side_effect = RuntimeError("TOP_SECRET browser crash")
|
|
self.assertEqual(pipeline.run_pipeline(self.args), 3)
|
|
self.assertEqual(lock_path.stat().st_ino, inode)
|
|
with lock_path.open("a") as lock:
|
|
fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB)
|
|
self.importer.assert_not_called()
|
|
self.assertNotIn("TOP_SECRET", self.stderr.getvalue())
|
|
|
|
def test_no_implicit_database(self):
|
|
with patch.dict(os.environ, {"PGUSER": "user", "PGDATABASE": "db"}, clear=True):
|
|
self.assertEqual(pipeline.run_pipeline(self.args), 3)
|
|
self.connect.assert_not_called()
|
|
self.collect.assert_not_called()
|
|
|
|
def test_lock_is_held_through_preflight_collection_and_import(self):
|
|
def assert_locked():
|
|
with (self.output / ".pipeline.lock").open("a") as lock:
|
|
with self.assertRaises(BlockingIOError):
|
|
fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB)
|
|
|
|
def collect(args):
|
|
assert_locked()
|
|
return self.fixture(args)
|
|
|
|
def import_run(conn, path):
|
|
assert_locked()
|
|
return self.summary
|
|
|
|
self.preflight.execute.side_effect = lambda query: assert_locked()
|
|
self.collect.side_effect = collect
|
|
self.importer.side_effect = import_run
|
|
self.assertEqual(pipeline.run_pipeline(self.args), 0)
|
|
with (self.output / ".pipeline.lock").open("a") as lock:
|
|
fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB)
|
|
|
|
def test_explicit_database_configuration(self):
|
|
for env, conninfo in (({"DATABASE_URL": "postgresql://configured"}, "postgresql://configured"),
|
|
({"PGSERVICE": "homelab"}, ""), ({"PGHOSTADDR": "192.0.2.1"}, "")):
|
|
with self.subTest(env=env), patch.dict(os.environ, env, clear=True):
|
|
self.preflight.reset_mock()
|
|
self.connect.side_effect = [self.preflight, self.import_connection]
|
|
self.assertEqual(pipeline.run_pipeline(self.args), 0)
|
|
self.connect.assert_called_with(conninfo, connect_timeout=10, autocommit=True)
|
|
|
|
def test_input_error_message_is_not_logged(self):
|
|
self.importer.side_effect = pipeline.InputError("TOP_SECRET")
|
|
self.assertEqual(pipeline.run_pipeline(self.args), 3)
|
|
self.assertNotIn("TOP_SECRET", self.stderr.getvalue() + json.dumps(self.state()))
|
|
|
|
def test_sidecar_storage_failure_prevents_import(self):
|
|
with patch.object(pipeline, "write_json", side_effect=OSError("TOP_SECRET")):
|
|
self.assertEqual(pipeline.run_pipeline(self.args), 3)
|
|
self.importer.assert_not_called()
|
|
self.assertTrue(self.raw.exists())
|
|
self.assertNotIn("TOP_SECRET", self.stderr.getvalue())
|
|
|
|
def test_main_delegates_parser_and_runner(self):
|
|
with patch.object(pipeline, "parse_collection_args", return_value=self.args) as parser:
|
|
with patch.object(pipeline, "run_pipeline", return_value=2) as run:
|
|
self.assertEqual(pipeline.main(["--search-only"]), 2)
|
|
parser.assert_called_once_with(["--search-only"], description=pipeline.__doc__)
|
|
run.assert_called_once_with(self.args)
|
|
|
|
def test_main_sanitizes_config_failures_and_handles_interrupt(self):
|
|
for error, code in ((ValueError("TOP_SECRET"), 3), (SystemExit(2), 3),
|
|
(KeyboardInterrupt(), 130), (SystemExit(0), 0)):
|
|
with self.subTest(error=type(error).__name__):
|
|
def fail(*args, **kwargs):
|
|
print("TOP_SECRET", file=pipeline.sys.stderr)
|
|
raise error
|
|
with patch.object(pipeline, "parse_collection_args", side_effect=fail):
|
|
self.assertEqual(pipeline.main([]), code)
|
|
self.assertNotIn("TOP_SECRET", self.stderr.getvalue())
|
|
|
|
def test_sigterm_is_handled_and_previous_handler_restored(self):
|
|
previous = signal.getsignal(signal.SIGTERM)
|
|
|
|
def terminated(args):
|
|
os.kill(os.getpid(), signal.SIGTERM)
|
|
|
|
self.collect.side_effect = terminated
|
|
with patch.object(pipeline, "parse_collection_args", return_value=self.args):
|
|
self.assertEqual(pipeline.main([]), 130)
|
|
self.importer.assert_not_called()
|
|
self.assertIs(signal.getsignal(signal.SIGTERM), previous)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|