Files
planet/scripts/harness/test_database_startup.py
linkong 60eefb19c9
Some checks are pending
ci / backend (push) Waiting to run
ci / frontend (push) Waiting to run
ci / delivery (push) Blocked by required conditions
release / images (push) Waiting to run
release: bump version to 0.74.7
2026-09-17 17:30:54 +08:00

639 lines
30 KiB
Python

"""Database lifecycle regressions without starting or changing host containers."""
import asyncio
from contextlib import redirect_stderr, redirect_stdout
import io
import json
from pathlib import Path
import re
import shlex
import socket
import subprocess
import sys
import tempfile
import unittest
from unittest.mock import AsyncMock, MagicMock, patch
from asyncpg import InvalidCatalogNameError, InvalidPasswordError
from sqlalchemy.engine import make_url
from sqlalchemy.exc import DBAPIError
ROOT = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(ROOT / "scripts"))
import check_database_connection as probe # noqa: E402
TEST_URL = make_url("postgresql+asyncpg://postgres:test-secret@localhost:5432/planet_db")
def shell_function(name: str) -> str:
source = (ROOT / "planet.sh").read_text() + "\n" + (ROOT / "scripts/lib/docker-bootstrap.zsh").read_text()
match = re.search(rf"^{name}\(\) \{{\n.*?^\}}", source, re.MULTILINE | re.DOTALL)
if match is None:
raise AssertionError(f"missing shell function: {name}")
return match.group()
def run_shell(functions: list[str], setup: str, action: str) -> subprocess.CompletedProcess[str]:
return subprocess.run(
["zsh", "-f"],
input="\n".join(["set -e", *(shell_function(name) for name in functions), setup, action]),
capture_output=True,
text=True,
timeout=10,
check=False,
)
class ComposeSelectionTests(unittest.TestCase):
def test_probe_timeout_does_not_fall_back_or_claim_compose_is_missing(self) -> None:
result = run_shell(
["compose_up"],
"""
compose_available() { return 124; }
log_error() { echo "$*"; }
compose_v1_available() { echo UNEXPECTED_FALLBACK; return 0; }
run_compose_command() { echo UNEXPECTED_COMMAND; }
report_missing_compose() { echo UNEXPECTED_MISSING; }
""",
"compose_up up -d postgres",
)
self.assertEqual(result.returncode, 124)
self.assertNotIn("UNEXPECTED", result.stdout)
def test_compose_timeout_keeps_live_output_and_passes_log_to_diagnostics(self) -> None:
with tempfile.TemporaryDirectory() as folder:
binary = Path(folder) / "docker"
binary.write_text("#!/bin/sh\necho 'postgres Pulling'\nsleep 30\n")
binary.chmod(0o755)
result = run_shell(
["run_compose_command", "run_docker_observed"],
f"""
PLANET_STATE_DIR={shlex.quote(folder)}
DOCKER_COMMAND_RUNNER={shlex.quote(str(ROOT / 'scripts/docker_command.py'))}
PLANET_COMPOSE_TIMEOUT=0.3
clear_wait_spinner() {{ :; }}
docker_failure_allows_mirror() {{ return 1; }}
docker_daemon_available() {{ return 0; }}
set_wait_detail() {{ echo "$*" >&2; }}
log_error() {{ echo "DIAGNOSTIC $2"; }}
""",
f"run_compose_command {shlex.quote(str(binary))} compose up -d postgres",
)
self.assertEqual(result.returncode, 124, result.stderr)
self.assertIn("postgres Pulling", result.stderr)
logs = list(Path(folder).glob("compose.*.log"))
self.assertEqual(len(logs), 1)
self.assertIn(f"DIAGNOSTIC {logs[0]}", result.stdout)
self.assertIn("PLANET_DOCKER_COMMAND_TIMEOUT", logs[0].read_text())
def test_operations_use_v1_only_when_v2_is_unavailable(self) -> None:
operations = (
("compose_up", "up -d postgres"),
("compose_supports_build", "build --help"),
("build_ai_provider_image", "build aiprovider"),
)
for function, arguments in operations:
for v2_available, v1_available in (
(True, True),
(True, False),
(False, True),
(False, False),
):
for command_status in (0, 1):
with self.subTest(
function=function,
v2=v2_available,
v1=v1_available,
command_status=command_status,
):
with tempfile.TemporaryDirectory() as directory:
calls_file = Path(directory) / "calls"
action = function
if function == "compose_up":
action += " " + arguments
result = run_shell(
[function],
f"""
compose_available() {{ return {int(not v2_available)}; }}
compose_v1_available() {{ return {int(not v1_available)}; }}
run_compose_command() {{ "$@"; }}
docker_probe() {{ "$@"; }}
docker() {{
echo "docker $*" >> '{calls_file}'
echo ORIGINAL_ERROR >&2
return {command_status}
}}
docker-compose() {{
echo "docker-compose $*" >> '{calls_file}'
return {command_status}
}}
run_ai_provider_build_command() {{
echo "$1 build aiprovider" >> '{calls_file}'
echo ORIGINAL_ERROR >&2
return {command_status}
}}
set_wait_detail() {{ :; }}
clear_wait_spinner() {{ :; }}
log_error() {{ echo "$*"; }}
log_note() {{ echo "$*"; }}
log_warn() {{ echo "$*"; }}
report_missing_compose() {{ echo MISSING_COMPOSE; return 1; }}
""",
f"if {action}; then exit 0; else exit $?; fi",
)
calls = (
calls_file.read_text().splitlines() if calls_file.exists() else []
)
command = "docker compose" if v2_available else "docker-compose"
available = v2_available or v1_available
self.assertEqual(calls, [f"{command} {arguments}"] if available else [])
self.assertEqual(result.returncode, command_status if available else 1)
self.assertNotIn("回退", result.stdout)
if available:
self.assertNotIn("MISSING_COMPOSE", result.stdout)
if v2_available and function != "compose_supports_build":
self.assertIn("ORIGINAL_ERROR", result.stderr)
class DatabaseLifecycleTests(unittest.TestCase):
def test_compose_error_exits_before_health_retry_or_restart(self) -> None:
result = run_shell(
["ensure_database_services_healthy"],
"""
DATABASE_START_MAX_RETRIES=3
start_database_services() { echo PORT_CONFLICT; return 1; }
wait_for_database_health() { echo UNEXPECTED_HEALTH; }
docker_probe() { echo UNEXPECTED_RESTART; }
animate_wait_spinner() { echo UNEXPECTED_RETRY; }
""",
"ensure_database_services_healthy",
)
self.assertEqual(result.returncode, 1)
self.assertEqual(result.stdout.strip(), "PORT_CONFLICT")
def test_failure_secondary_line_survives_outer_error_and_cleanup_notice(self) -> None:
result = run_shell(
["log_error", "cleanup_failed_start"],
"""
WAIT_SPINNER_MESSAGE='启动数据库'
LAST_COMMAND_ERROR='Bind for 127.0.0.1:5432 failed: port is already allocated'
VERBOSE=0
log_line() { echo "$1 $3"; }
stop_wait_session() { :; }
clear_wait_spinner() { echo UNEXPECTED_CLEAR; }
log_warn() { echo UNEXPECTED_WARNING; }
START_RUN_ACTIVE=1
START_RUN_COMPLETED=0
STARTED_MOTION_AGENT_THIS_RUN=0
STARTED_FRONTEND_THIS_RUN=0
STARTED_BACKEND_THIS_RUN=0
""",
"log_error 'Docker Compose 执行失败'\n"
"log_error '后端数据库连接检查失败'\ncleanup_failed_start",
)
self.assertEqual(result.returncode, 0, result.stderr)
self.assertEqual(result.stdout.splitlines(), [
"fail 启动数据库",
" Bind for 127.0.0.1:5432 failed: port is already allocated",
])
def test_ai_start_uses_host_readiness_without_waiting_for_docker_probe_schedule(self) -> None:
for recreate in (0, 1):
with self.subTest(recreate=recreate):
result = run_shell(
["start_ai_provider_service"],
f"""
AI_PROVIDER_RECREATE_REQUIRED={recreate}
AI_PROVIDER_START_MAX_RETRIES=1
AI_PROVIDER_HEALTH_CHECK_ATTEMPTS=10
AI_PROVIDER_HEALTH_CHECK_INTERVAL=2
AI_PROVIDER_CONTAINER_NAME=planet_aiprovider
for fn in set_wait_detail write_ai_provider_runtime_env_file \
ensure_ai_provider_image_current recreate_ai_provider_container docker_probe; do
functions[$fn]='return 0'
done
wait_for_container_health() {{ echo WAIT_FOR_DOCKER_SCHEDULE; return 1; }}
wait_for_http() {{ echo "HOST_CHECK $1"; return 0; }}
log_error() {{ :; }}
""",
"start_ai_provider_service 18010",
)
self.assertEqual(result.returncode, 0, result.stdout + result.stderr)
self.assertIn("HOST_CHECK http://localhost:18010/health", result.stdout)
self.assertNotIn("WAIT_FOR_DOCKER_SCHEDULE", result.stdout)
def test_fatal_backend_startup_does_not_launch_three_identical_attempts(self) -> None:
result = run_shell(
["start_backend_with_retry"],
f"""
SCRIPT_DIR={ROOT}
BACKEND_LOG_FILE=/dev/null
BACKEND_PID_FILE=/dev/null
BACKEND_MAX_RETRIES=3
BACKEND_HEALTH_CHECK_ATTEMPTS=60
BACKEND_HEALTH_CHECK_INTERVAL=2
VERBOSE=0
cleanup_backend_processes() {{ :; }}
start_detached_command() {{ echo STARTED >&2; echo 42; }}
write_pid_file() {{ :; }}
wait_for_http() {{ return 2; }}
kill() {{ :; }}
backend_log_indicates_port_conflict() {{ return 1; }}
animate_wait_spinner() {{ echo WAITED; }}
log_error() {{ :; }}
""",
"start_backend_with_retry 8000",
)
self.assertEqual(result.returncode, 2, result.stdout + result.stderr)
self.assertEqual(result.stderr.count("STARTED"), 1)
self.assertNotIn("WAITED", result.stdout)
def test_start_keeps_stopped_containers_available_for_reuse(self) -> None:
result = run_shell(
["start"],
"""
BACKEND_PORT=8000
FRONTEND_PORT=3000
AI_PROVIDER_PORT=8010
BACKEND_PORT_REQUESTED=0
FRONTEND_PORT_REQUESTED=0
FRONTEND_LAN_ENABLED=0
MOTION_AGENT_REQUESTED=0
for fn in parse_service_args prepare_allow_lan_public_ports print_splash \
start_backend_service start_frontend_service write_port_state \
verify_allow_lan_access log_success log_note; do
functions[$fn]='return 0'
done
current_frontend_scheme() { echo http; }
cleanup_exit_containers() { echo DELETED_STOPPED_CONTAINERS; }
""",
"start --non-motion-agent",
)
self.assertEqual(result.returncode, 0, result.stderr)
self.assertNotIn("DELETED_STOPPED_CONTAINERS", result.stdout)
def test_start_checks_database_before_preparing_ai_and_starting_backend(self) -> None:
for connection_status in (0, 1):
with self.subTest(connection_status=connection_status):
result = run_shell(
["start_backend_service"],
f"""
BACKEND_MAX_RETRIES=3
BACKEND_LOG_FILE=/dev/null
for fn in set_wait_detail start_wait_session stop_wait_session log_success log_note log_error; do
functions[$fn]='return 0'
done
ensure_uv_backend_deps() {{ echo DEPS; }}
ensure_database_services_healthy() {{ echo CONTAINERS; }}
verify_backend_database_connection() {{ echo DB_CHECK; return {connection_status}; }}
ai_provider_service_healthy() {{ return 1; }}
start_ai_provider_service() {{ echo AI_STARTED; }}
start_backend_with_retry() {{ echo BACKEND_STARTED; }}
""",
"start_backend_service 8000 0 8010",
)
self.assertIn("DB_CHECK", result.stdout)
if connection_status:
self.assertNotEqual(result.returncode, 0)
self.assertNotIn("AI_STARTED", result.stdout)
self.assertNotIn("BACKEND_STARTED", result.stdout)
else:
self.assertEqual(result.returncode, 0, result.stderr)
self.assertLess(result.stdout.index("DB_CHECK"), result.stdout.index("AI_STARTED"))
self.assertEqual(result.stdout.count("DEPS"), 1)
def test_backend_fatal_startup_log_stops_waiting_while_reloader_is_alive(self) -> None:
with tempfile.TemporaryDirectory() as directory:
log_file = Path(directory) / "backend.log"
log_file.write_text("Application startup failed. Exiting.\n")
result = run_shell(
["wait_for_http", "backend_startup_failed", "log_matches"],
f"""
BACKEND_PID=$$
BACKEND_LOG_FILE={log_file}
set_wait_detail() {{ :; }}
clear_wait_spinner() {{ :; }}
http_ok() {{ echo HTTP_PROBE; return 1; }}
animate_wait_spinner() {{ echo WAIT; }}
""",
"wait_for_http http://localhost:8000/health 60 2 backend backend_startup_failed",
)
self.assertEqual(result.returncode, 2, result.stdout + result.stderr)
self.assertNotIn("WAIT", result.stdout)
self.assertNotIn("HTTP_PROBE", result.stdout)
def test_readiness_without_failure_probe_still_waits_for_slow_healthy_services(self) -> None:
result = run_shell(
["wait_for_http"],
"""
probes=0
set_wait_detail() { :; }
clear_wait_spinner() { :; }
finish_wait_spinner() { :; }
format_wait_elapsed_seconds() { echo 0; }
http_ok() { probes=$((probes + 1)); [ $probes -eq 3 ]; }
animate_wait_spinner() { echo WAIT; }
""",
"wait_for_http http://localhost:8010/health 10 2 provider",
)
self.assertEqual(result.returncode, 0, result.stderr)
self.assertEqual(result.stdout.count("WAIT"), 2)
def test_backend_worker_import_failure_is_distinct_from_optional_collector_errors(self) -> None:
for contents, expected_status in (
("Process SpawnProcess-1:\nModuleNotFoundError: missing_app\n", 0),
("Collector unavailable\nImportError: optional_dependency\n", 1),
):
with self.subTest(contents=contents), tempfile.TemporaryDirectory() as directory:
log_file = Path(directory) / "backend.log"
log_file.write_text(contents)
result = run_shell(
["backend_startup_failed", "log_matches"],
f"BACKEND_PID=$$\nBACKEND_LOG_FILE={log_file}",
"backend_startup_failed",
)
self.assertEqual(result.returncode, expected_status, result.stderr)
def test_existing_container_gets_current_compose_configuration(self) -> None:
for function in ("start_database_services", "start_postgres_service"):
with self.subTest(function=function):
result = run_shell(
[function],
"""
mapped=0
prepare_database_ports() { return 0; }
docker() { return 0; }
compose_up() { mapped=1; }
""",
f"{function}\n[ $mapped -eq 1 ]",
)
self.assertEqual(result.returncode, 0, result.stdout + result.stderr)
def test_port_conflict_is_visible_and_cannot_fall_back_to_old_container(self) -> None:
result = run_shell(
["start_database_services"],
"""
docker() { return 0; }
prepare_database_ports() { return 0; }
compose_up() { echo 'port is already allocated' >&2; return 1; }
""",
"start_database_services",
)
self.assertNotEqual(result.returncode, 0)
self.assertIn("port is already allocated", result.stderr)
def test_installed_compose_failure_is_not_reported_as_missing_compose(self) -> None:
result = run_shell(
["compose_up"],
"""
compose_available() { return 0; }
compose_v1_available() { return 1; }
run_compose_command() { "$@"; }
docker() { echo 'address already in use' >&2; return 1; }
set_wait_detail() { :; }
log_warn() { echo "$*"; }
log_error() { echo "$*"; }
report_missing_compose() { echo MISSING_COMPOSE; return 1; }
""",
"compose_up up -d postgres",
)
self.assertNotEqual(result.returncode, 0)
self.assertIn("address already in use", result.stderr)
self.assertNotIn("MISSING_COMPOSE", result.stdout)
def run_init(self, connection_status: int) -> subprocess.CompletedProcess[str]:
return run_shell(
["init"],
f"""
for fn in parse_service_args guard_init_when_services_running print_splash \
start_wait_session stop_wait_session ensure_docker_runtime ensure_python_runtime \
sync_python_deps sync_frontend_deps ensure_planet_env_files \
prepare_motion_agent_host_dependencies set_wait_detail; do
functions[$fn]='return 0'
done
log_success() {{ echo "$*"; }}
log_error() {{ echo "$*"; }}
log_step() {{ :; }}
log_note() {{ :; }}
ensure_database_services_healthy() {{ echo CONTAINERS_HEALTHY; }}
verify_backend_database_connection() {{ echo CONNECTION_CHECK; return {connection_status}; }}
run_command_with_spinner() {{ shift; "$@"; }}
initialize_backend_database() {{ echo SCHEMA_INITIALIZED; }}
""",
"init --non-motion-agent",
)
def test_healthy_containers_do_not_allow_schema_changes_on_bad_connection(self) -> None:
result = self.run_init(1)
self.assertNotEqual(result.returncode, 0, result.stdout)
self.assertIn("CONNECTION_CHECK", result.stdout)
self.assertNotIn("SCHEMA_INITIALIZED", result.stdout)
self.assertNotIn("数据库服务已就绪", result.stdout)
def test_real_connection_is_checked_before_reporting_ready_and_creating_tables(self) -> None:
result = self.run_init(0)
self.assertEqual(result.returncode, 0, result.stderr)
self.assertIn("CONNECTION_CHECK", result.stdout)
self.assertLess(
result.stdout.index("CONNECTION_CHECK"), result.stdout.index("数据库服务已就绪")
)
self.assertLess(
result.stdout.index("数据库服务已就绪"), result.stdout.index("SCHEMA_INITIALIZED")
)
def test_only_missing_ports_trigger_one_volume_preserving_recreation(self) -> None:
for first_status, retry_status, expected_status, expected_recreates in (
(0, 0, 0, 0),
(1, 0, 1, 0),
(2, 0, 0, 1),
(2, 2, 2, 1),
):
with self.subTest(first_status=first_status, retry_status=retry_status):
result = run_shell(
["verify_backend_database_connection"],
f"""
checks=0
mktemp() {{ echo /dev/null; }}
set_wait_detail() {{ :; }}
log_warn() {{ :; }}
compose_up() {{ echo "COMPOSE $*"; }}
wait_for_postgres_health() {{ return 0; }}
run_docker_observed() {{
checks=$((checks + 1))
if [ $checks -eq 1 ]; then return {first_status}; fi
return {retry_status}
}}
""",
"verify_backend_database_connection",
)
self.assertEqual(result.returncode, expected_status, result.stderr)
self.assertEqual(result.stdout.count("COMPOSE"), expected_recreates)
if expected_recreates:
self.assertIn("up -d --no-deps --force-recreate postgres", result.stdout)
self.assertNotIn("--renew-anon-volumes", result.stdout)
class DatabaseProbeTests(unittest.TestCase):
def test_recently_closed_connection_is_not_a_port_conflict(self) -> None:
with socket.socket() as server, socket.socket() as client:
server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
server.bind(("127.0.0.1", 0))
port = server.getsockname()[1]
server.listen()
client.connect(("127.0.0.1", port))
connection, _ = server.accept()
connection.close()
client.recv(1)
result = subprocess.CompletedProcess([], 0, "", "")
with patch.object(probe.subprocess, "run", return_value=result):
probe.check_port_conflicts([("PostgreSQL", "planet_postgres", port)])
def test_actual_host_listener_is_still_a_conflict(self) -> None:
with socket.socket() as server:
server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
server.bind(("127.0.0.1", 0))
server.listen()
port = server.getsockname()[1]
result = subprocess.CompletedProcess([], 0, "", "")
with patch.object(probe.subprocess, "run", return_value=result):
with self.assertRaises(probe.DatabaseReadinessError):
probe.check_port_conflicts([("PostgreSQL", "planet_postgres", port)])
def test_other_project_database_and_redis_ports_are_identified_before_recreate(self) -> None:
result = subprocess.CompletedProcess([], 0,
"happyrouter-db-1\t127.0.0.1:5432->5432/tcp\n"
"happyrouter-redis-1\t127.0.0.1:6379->6379/tcp\n", "")
for service, container, port, owner in (
("PostgreSQL", "planet_postgres", 5432, "happyrouter-db-1"),
("Redis", "planet_redis", 6379, "happyrouter-redis-1"),
):
with patch.object(probe.subprocess, "run", return_value=result):
with self.assertRaises(probe.DatabaseReadinessError) as error:
probe.check_port_conflicts([(service, container, port)])
self.assertIn(owner, str(error.exception))
self.assertIn("PLANET_DB_PORT_OCCUPIED", str(error.exception))
def test_own_published_port_is_reusable_and_other_port_is_not_a_conflict(self) -> None:
result = subprocess.CompletedProcess([], 0,
"planet_postgres\t127.0.0.1:15432->5432/tcp\n"
"happyrouter-db-1\t127.0.0.1:5432->5432/tcp\n", "")
with patch.object(probe.subprocess, "run", return_value=result), \
patch.object(probe.socket, "socket") as bind:
probe.check_port_conflicts([("PostgreSQL", "planet_postgres", 15432)])
bind.assert_not_called()
def docker_result(
self, ports: dict | None, mode: str = "bridge"
) -> subprocess.CompletedProcess:
return subprocess.CompletedProcess([], 0, f"{mode}\n{json.dumps(ports)}\n", "")
def test_missing_port_blocks_sql_even_if_container_is_healthy(self) -> None:
output = io.StringIO()
with (
patch.object(probe, "backend_database_url", return_value=TEST_URL),
patch.object(probe.subprocess, "run", return_value=self.docker_result({})),
patch.object(probe, "check_connection", new_callable=AsyncMock) as connect,
redirect_stderr(output),
redirect_stdout(output),
):
self.assertEqual(probe.main(), 2)
connect.assert_not_called()
self.assertIn("未向宿主机发布", output.getvalue())
self.assertNotIn("test-secret", output.getvalue())
def test_actual_backend_port_must_match_published_port(self) -> None:
result = self.docker_result({"5432/tcp": [{"HostIp": "0.0.0.0", "HostPort": "15432"}]})
with patch.object(probe.subprocess, "run", return_value=result):
with self.assertRaises(probe.DatabaseReadinessError):
probe.check_published_port(TEST_URL)
probe.check_published_port(TEST_URL.set(port=15432))
def test_host_networking_does_not_require_a_port_mapping(self) -> None:
with patch.object(probe.subprocess, "run", return_value=self.docker_result(None, "host")):
probe.check_published_port(TEST_URL)
def test_external_database_does_not_require_a_local_mapping(self) -> None:
with patch.object(probe.subprocess, "run") as inspect:
probe.check_published_port(TEST_URL.set(host="configured-db.example"))
inspect.assert_not_called()
def test_docker_failure_does_not_echo_raw_output(self) -> None:
result = subprocess.CompletedProcess([], 1, "", "sensitive test-secret")
with patch.object(probe.subprocess, "run", return_value=result):
with self.assertRaises(probe.DatabaseReadinessError) as error:
probe.check_published_port(TEST_URL)
self.assertIn("Docker context", str(error.exception))
self.assertNotIn("test-secret", str(error.exception))
def test_authentication_error_is_actionable_without_leaking_dsn_or_query_secrets(self) -> None:
output = io.StringIO()
url = TEST_URL.update_query_dict({"sslpassword": "query-secret"})
error = DBAPIError(
None, None, InvalidPasswordError(url.render_as_string(hide_password=False))
)
with (
patch.object(probe, "backend_database_url", return_value=url),
patch.object(probe, "check_published_port"),
patch.object(probe, "check_connection", new_callable=AsyncMock, side_effect=error),
redirect_stderr(output),
redirect_stdout(output),
):
self.assertEqual(probe.main(), 1)
self.assertIn("认证失败", output.getvalue())
self.assertNotIn("test-secret", output.getvalue())
self.assertNotIn("query-secret", output.getvalue())
self.assertNotIn("Traceback", output.getvalue())
def test_missing_database_and_network_errors_have_distinct_diagnostics(self) -> None:
self.assertIn("数据库不存在", probe.connection_diagnostic(InvalidCatalogNameError()))
for error in (ConnectionRefusedError(), TimeoutError()):
self.assertIn("地址不可达", probe.connection_diagnostic(error))
def test_invalid_config_never_echoes_validation_input(self) -> None:
output = io.StringIO()
with (
patch.object(probe, "backend_database_url", side_effect=ValueError("test-secret")),
redirect_stderr(output),
):
self.assertEqual(probe.main(), 1)
self.assertNotIn("test-secret", output.getvalue())
class DatabaseConnectionTests(unittest.IsolatedAsyncioTestCase):
async def test_connection_deadline_stops_a_stalled_probe(self) -> None:
async def stall() -> None:
await asyncio.sleep(1)
engine = MagicMock()
engine.dispose = AsyncMock()
engine.connect.return_value.__aenter__.side_effect = stall
with (
patch.object(probe, "create_async_engine", return_value=engine),
patch.object(probe, "CONNECT_TIMEOUT_SECONDS", 0.001),
):
with self.assertRaises(TimeoutError):
await probe.check_connection(TEST_URL)
engine.dispose.assert_awaited_once()
async def test_probe_only_selects_and_always_disposes_the_engine(self) -> None:
for error in (None, InvalidPasswordError("test-secret")):
with self.subTest(error=type(error).__name__):
engine = MagicMock()
engine.dispose = AsyncMock()
connection = AsyncMock()
engine.connect.return_value.__aenter__.return_value = connection
connection.execute.side_effect = error
with patch.object(probe, "create_async_engine", return_value=engine):
if error:
with self.assertRaises(InvalidPasswordError):
await probe.check_connection(TEST_URL)
else:
await probe.check_connection(TEST_URL)
self.assertEqual(str(connection.execute.call_args.args[0]), "SELECT 1")
engine.dispose.assert_awaited_once()
if __name__ == "__main__":
unittest.main()