237 lines
11 KiB
Python
237 lines
11 KiB
Python
"""Automatic proxy routing and rollback tests; never restart the host Docker daemon."""
|
|
|
|
import json
|
|
from pathlib import Path
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
import unittest
|
|
from unittest.mock import patch
|
|
|
|
from test_database_startup import run_shell
|
|
|
|
ROOT = Path(__file__).resolve().parents[2]
|
|
sys.path.insert(0, str(ROOT / "scripts"))
|
|
|
|
import docker_proxy as proxy # noqa: E402
|
|
|
|
URLS = ["https://registry.example/v2/"]
|
|
PROXIES = {
|
|
"http-proxy": "http://localhost:1234",
|
|
"https-proxy": "http://localhost:1234",
|
|
"no-proxy": "localhost,127.0.0.1,::1",
|
|
}
|
|
|
|
|
|
class DockerProxyTests(unittest.TestCase):
|
|
def setUp(self) -> None:
|
|
self.temporary = tempfile.TemporaryDirectory(prefix="planet-proxy-test-")
|
|
self.addCleanup(self.temporary.cleanup)
|
|
self.config = Path(self.temporary.name) / "daemon.json"
|
|
self.state = Path(self.temporary.name) / "proxy-state.json"
|
|
|
|
def test_no_host_proxy_uses_direct_without_adding_configuration(self) -> None:
|
|
with patch.object(proxy, "reachable", return_value=True) as probe:
|
|
plan = proxy.select_route({}, URLS)
|
|
self.assertEqual(plan, {"desired": {}, "mode": "direct", "reachable": True})
|
|
probe.assert_called_once_with(URLS, "")
|
|
plan["current"] = {}
|
|
with (
|
|
patch.object(proxy, "daemon_proxies", return_value={}),
|
|
patch.object(proxy, "run") as run,
|
|
):
|
|
proxy.apply_plan(plan, self.config, self.state)
|
|
run.assert_not_called()
|
|
self.assertFalse(self.config.exists())
|
|
|
|
def test_proxy_is_selected_only_after_registry_probe(self) -> None:
|
|
with patch.object(proxy, "reachable", return_value=True) as probe:
|
|
plan = proxy.select_route({"https_proxy": PROXIES["https-proxy"]}, URLS)
|
|
self.assertEqual(plan["desired"], PROXIES)
|
|
probe.assert_called_once_with(URLS, PROXIES["https-proxy"], PROXIES["no-proxy"])
|
|
|
|
def test_dead_proxy_falls_back_to_direct(self) -> None:
|
|
with patch.object(proxy, "reachable", side_effect=[False, True]) as probe:
|
|
plan = proxy.select_route({"HTTPS_PROXY": PROXIES["https-proxy"]}, URLS)
|
|
self.assertEqual(plan["desired"], {})
|
|
self.assertTrue(plan["reachable"])
|
|
self.assertEqual(probe.call_count, 2)
|
|
|
|
def test_both_routes_failing_is_not_reported_as_connected(self) -> None:
|
|
with patch.object(proxy, "reachable", return_value=False):
|
|
plan = proxy.select_route({"HTTP_PROXY": PROXIES["http-proxy"]}, URLS)
|
|
self.assertEqual(plan["mode"], "direct")
|
|
self.assertFalse(plan["reachable"])
|
|
|
|
def test_no_proxy_rules_are_preserved_in_probe_and_config(self) -> None:
|
|
with patch.object(proxy, "reachable", return_value=True) as probe:
|
|
plan = proxy.select_route(
|
|
{"HTTPS_PROXY": PROXIES["https-proxy"], "NO_PROXY": "*"}, URLS
|
|
)
|
|
self.assertEqual(plan["desired"]["no-proxy"], "*")
|
|
probe.assert_called_once_with(URLS, PROXIES["https-proxy"], "*")
|
|
|
|
def test_actual_image_registry_overrides_are_probed(self) -> None:
|
|
urls = proxy.registry_urls(
|
|
{
|
|
"PYTHON_IMAGE": "internal.example:5000/python:3",
|
|
"UV_IMAGE": "internal.example:5000/uv:latest",
|
|
}
|
|
)
|
|
self.assertEqual(urls, ["https://internal.example:5000/v2/"])
|
|
|
|
def test_only_actual_builds_prepare_proxy_and_failed_preparation_stops_build(self) -> None:
|
|
for no_build, fingerprint, proxy_status, expected_calls in (
|
|
(1, "old", 0, []),
|
|
(0, "new", 0, []),
|
|
(0, "old", 0, ["PROXY_CHECK", "BUILD"]),
|
|
(0, "old", 1, ["PROXY_CHECK"]),
|
|
):
|
|
with self.subTest(no_build=no_build, fingerprint=fingerprint, status=proxy_status):
|
|
result = run_shell(
|
|
["ensure_ai_provider_image_current"],
|
|
f"""
|
|
AI_PROVIDER_NO_BUILD={no_build}
|
|
AI_PROVIDER_BUILD_LOG_FILE=/dev/null
|
|
prepare_uv_build_config() {{ :; }}
|
|
compute_ai_provider_build_fingerprint() {{ echo new; }}
|
|
read_ai_provider_build_stamp() {{ echo old; }}
|
|
ai_provider_image_exists() {{ return 0; }}
|
|
read_ai_provider_image_fingerprint() {{ echo {fingerprint}; }}
|
|
write_ai_provider_build_stamp() {{ :; }}
|
|
set_wait_detail() {{ :; }}
|
|
log_note() {{ :; }}
|
|
log_success() {{ :; }}
|
|
compose_supports_build() {{ return 0; }}
|
|
prepare_docker_build_proxy() {{ echo PROXY_CHECK; return {proxy_status}; }}
|
|
build_ai_provider_image() {{ echo BUILD; }}
|
|
""",
|
|
"ensure_ai_provider_image_current",
|
|
)
|
|
self.assertEqual(result.stdout.splitlines(), expected_calls)
|
|
self.assertEqual(result.returncode, proxy_status)
|
|
|
|
def fake_run(self, args: list[str], **kwargs: object) -> subprocess.CompletedProcess[str]:
|
|
output = "postgres-id\nredis-id\n" if args == ["docker", "ps", "-q"] else ""
|
|
return subprocess.CompletedProcess(args, 0, output, "")
|
|
|
|
def test_add_proxy_preserves_other_settings_and_restores_containers(self) -> None:
|
|
original = b'{"log-driver": "local"}\n'
|
|
self.config.write_bytes(original)
|
|
with (
|
|
patch.object(proxy, "daemon_proxies", side_effect=[{}, PROXIES]),
|
|
patch.object(proxy, "run", side_effect=self.fake_run) as run,
|
|
):
|
|
proxy.apply_plan({"current": {}, "desired": PROXIES}, self.config, self.state)
|
|
self.assertEqual(
|
|
json.loads(self.config.read_text()), {"log-driver": "local", "proxies": PROXIES}
|
|
)
|
|
self.assertEqual(json.loads(self.state.read_text())["managed"], PROXIES)
|
|
self.assertEqual(self.config.with_suffix(".json.planet-proxy.bak").read_bytes(), original)
|
|
self.assertEqual(self.state.stat().st_mode & 0o777, 0o600)
|
|
run.assert_any_call(["docker", "start", "postgres-id", "redis-id"])
|
|
|
|
def test_remove_owned_proxy_when_host_proxy_disappears(self) -> None:
|
|
self.config.write_text(json.dumps({"proxies": PROXIES, "log-driver": "local"}))
|
|
self.state.write_text(json.dumps({"managed": PROXIES}))
|
|
with (
|
|
patch.object(proxy, "daemon_proxies", side_effect=[PROXIES, {}]),
|
|
patch.object(proxy, "run", side_effect=self.fake_run),
|
|
):
|
|
proxy.apply_plan({"current": PROXIES, "desired": {}}, self.config, self.state)
|
|
self.assertEqual(json.loads(self.config.read_text()), {"log-driver": "local"})
|
|
|
|
def test_unchanged_proxy_does_not_restart_docker(self) -> None:
|
|
with (
|
|
patch.object(proxy, "daemon_proxies", return_value=PROXIES),
|
|
patch.object(proxy, "run") as run,
|
|
):
|
|
proxy.apply_plan({"current": PROXIES, "desired": PROXIES}, self.config, self.state)
|
|
run.assert_not_called()
|
|
|
|
def test_user_managed_proxy_is_not_overwritten(self) -> None:
|
|
original = json.dumps({"proxies": PROXIES})
|
|
self.config.write_text(original)
|
|
with (
|
|
patch.object(proxy, "daemon_proxies", return_value=PROXIES),
|
|
patch.object(proxy, "run") as run,
|
|
self.assertRaisesRegex(proxy.ProxyError, "PLANET_PROXY_EXTERNAL"),
|
|
):
|
|
proxy.apply_plan({"current": PROXIES, "desired": {}}, self.config, self.state)
|
|
run.assert_not_called()
|
|
self.assertEqual(self.config.read_text(), original)
|
|
|
|
def test_external_edit_of_managed_config_is_not_overwritten(self) -> None:
|
|
changed = {**PROXIES, "https-proxy": "http://different.example:8080"}
|
|
self.config.write_text(json.dumps({"proxies": changed}))
|
|
self.state.write_text(json.dumps({"managed": PROXIES}))
|
|
with (
|
|
patch.object(proxy, "daemon_proxies", return_value=changed),
|
|
self.assertRaisesRegex(proxy.ProxyError, "PLANET_PROXY_EXTERNAL"),
|
|
):
|
|
proxy.apply_plan({"current": changed, "desired": {}}, self.config, self.state)
|
|
|
|
def test_restart_failure_restores_original_config_and_containers(self) -> None:
|
|
original = b'{"log-driver": "local"}\n'
|
|
self.config.write_bytes(original)
|
|
restarts = 0
|
|
|
|
def fail_first_restart(args: list[str]) -> subprocess.CompletedProcess[str]:
|
|
nonlocal restarts
|
|
if args == ["systemctl", "restart", "docker"]:
|
|
restarts += 1
|
|
if restarts == 1:
|
|
raise subprocess.CalledProcessError(1, args)
|
|
return self.fake_run(args)
|
|
|
|
with (
|
|
patch.object(proxy, "daemon_proxies", return_value={}),
|
|
patch.object(proxy, "run", side_effect=fail_first_restart) as run,
|
|
self.assertRaisesRegex(proxy.ProxyError, "PLANET_PROXY_CONFIG_FAILED"),
|
|
):
|
|
proxy.apply_plan({"current": {}, "desired": PROXIES}, self.config, self.state)
|
|
self.assertEqual(self.config.read_bytes(), original)
|
|
self.assertFalse(self.state.exists())
|
|
self.assertEqual(restarts, 2)
|
|
run.assert_any_call(["docker", "start", "postgres-id", "redis-id"])
|
|
|
|
def test_probe_does_not_put_proxy_credentials_in_command_arguments(self) -> None:
|
|
with patch.object(
|
|
proxy.subprocess, "run", return_value=subprocess.CompletedProcess([], 0, "401")
|
|
) as run:
|
|
self.assertTrue(proxy.probe_url(URLS[0], "http://user:secret@proxy.example:1234"))
|
|
self.assertNotIn("secret", " ".join(run.call_args.args[0]))
|
|
self.assertIn("secret", run.call_args.kwargs["env"]["https_proxy"])
|
|
|
|
def test_registry_rejection_does_not_mark_the_proxy_unreachable(self) -> None:
|
|
for status in ("401", "403", "429"):
|
|
with (
|
|
self.subTest(status=status),
|
|
patch.object(
|
|
proxy.subprocess, "run", return_value=subprocess.CompletedProcess([], 0, status)
|
|
),
|
|
):
|
|
self.assertTrue(proxy.probe_url(URLS[0], PROXIES["https-proxy"]))
|
|
|
|
def test_validation_failure_restores_config_without_restarting_docker(self) -> None:
|
|
original = b'{"log-driver": "local"}\n'
|
|
self.config.write_bytes(original)
|
|
|
|
def fail_validation(args: list[str]) -> subprocess.CompletedProcess[str]:
|
|
if args[0] == "dockerd":
|
|
raise subprocess.CalledProcessError(1, args)
|
|
return self.fake_run(args)
|
|
|
|
with (
|
|
patch.object(proxy, "daemon_proxies", return_value={}),
|
|
patch.object(proxy, "run", side_effect=fail_validation) as run,
|
|
self.assertRaisesRegex(proxy.ProxyError, "PLANET_PROXY_CONFIG_FAILED"),
|
|
):
|
|
proxy.apply_plan({"current": {}, "desired": PROXIES}, self.config, self.state)
|
|
self.assertEqual(self.config.read_bytes(), original)
|
|
self.assertFalse(any(call.args[0][0] == "systemctl" for call in run.call_args_list))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|