diff --git a/main.py b/main.py index 22ffbd2..a0d2f25 100644 --- a/main.py +++ b/main.py @@ -25,6 +25,9 @@ REGISTRY_CUTOVER_YEAR = 2026 REGISTRY_CUTOVER_MONTH = 4 +# Matches the Connect server's command line, for pkill -f +CONNECT_PROCESS_PATTERN = "^/opt/rstudio-connect/bin/connect( |$)" + DESCRIPTION = """\ Starts Posit Connect in Docker and waits until it is ready. @@ -274,6 +277,28 @@ def get_docker_tag(version: str) -> tuple[str, str]: return (LEGACY_IMAGE, version) +def stop_connect_container(container) -> None: + """ + Stop a Connect container promptly. + + Newer images run startup.sh as PID 1, and a shell running as PID 1 + ignores SIGTERM, so container.stop() waits out its 10 second grace + period and then kills the container. Signaling the Connect process + directly lets it shut down cleanly in under a second, and the container + exits with it. If that doesn't work, fall back to a short stop timeout. + """ + try: + exit_code, _ = container.exec_run( + ["pkill", "-TERM", "-f", CONNECT_PROCESS_PATTERN] + ) + if exit_code == 0: + container.wait(timeout=10) + return + except Exception: + pass # e.g. no pkill in the image, or Connect didn't exit in time + container.stop(timeout=2) + + def main() -> int: """ Main entry point for the with-connect CLI tool. @@ -297,7 +322,7 @@ def main() -> int: client = docker.from_env() try: container = client.containers.get(container_id) - container.stop() + stop_connect_container(container) print(f"Stopped container {container_id}", file=sys.stderr) return 0 except docker.errors.NotFound: @@ -417,7 +442,7 @@ def main() -> int: return exit_code finally: if stop_container: - container.stop() + stop_connect_container(container) def is_port_open(host: str, port: int, timeout: float = 30.0) -> bool: @@ -476,7 +501,6 @@ def wait_for_http_server( if "Unable to obtain a valid license" in logs: print("\nContainer logs:", file=sys.stderr) print(logs, file=sys.stderr) - container.stop() raise RuntimeError( "Unable to obtain a valid license. Your Posit Connect license may be expired or invalid. Please check your license file." ) diff --git a/test_main.py b/test_main.py index 41525dc..6da2fc1 100644 --- a/test_main.py +++ b/test_main.py @@ -13,6 +13,7 @@ import docker import main +import requests def test_license_file_not_exists(): @@ -75,7 +76,8 @@ def test_invalid_license_detection(): except RuntimeError as e: assert "Unable to obtain a valid license" in str(e) assert "expired or invalid" in str(e) - mock_container.stop.assert_called_once() + # main() stops the container when this raises + mock_container.stop.assert_not_called() def test_valid_license_http_server_starts(): @@ -346,6 +348,83 @@ def test_stop_nonexistent_container(): assert "Container not found" in result.stderr +def test_stop_connect_container_signals_connect(): + container = MagicMock() + container.exec_run.return_value = (0, b"") + + main.stop_connect_container(container) + + container.exec_run.assert_called_once_with( + ["pkill", "-TERM", "-f", main.CONNECT_PROCESS_PATTERN] + ) + container.wait.assert_called_once_with(timeout=10) + container.stop.assert_not_called() + + +def test_stop_mode_signals_connect(): + client = MagicMock() + container = client.containers.get.return_value + container.exec_run.return_value = (0, b"") + with ( + patch.object(sys, "argv", ["with-connect", "--stop", "abc123"]), + patch.object(main.docker, "from_env", return_value=client), + redirect_stderr(io.StringIO()), + ): + assert main.main() == 0 + + container.exec_run.assert_called_once_with( + ["pkill", "-TERM", "-f", main.CONNECT_PROCESS_PATTERN] + ) + + +def test_command_mode_signals_connect_after_command(): + client = MagicMock() + container = client.containers.run.return_value + container.exec_run.return_value = (0, b"") + with tempfile.NamedTemporaryFile(suffix=".lic") as license_file: + argv = ["with-connect", "--license", license_file.name, "--", "true"] + with ( + patch.object(sys, "argv", argv), + patch.object(main.docker, "from_env", return_value=client), + patch.object(main, "ensure_image"), + patch.object(main, "is_port_open", return_value=True), + patch.object(main, "wait_for_http_server", return_value=True), + patch.object(main, "get_api_key", return_value="key"), + patch.object(main.subprocess, "run", return_value=Mock(returncode=0)), + redirect_stderr(io.StringIO()), + ): + assert main.main() == 0 + + container.exec_run.assert_called_once_with( + ["pkill", "-TERM", "-f", main.CONNECT_PROCESS_PATTERN] + ) + + +def test_stop_connect_container_falls_back_when_connect_not_signaled(): + # pkill exits 1 when nothing matches and 127 when it isn't installed + for exit_code in (1, 127): + container = MagicMock() + container.exec_run.return_value = (exit_code, b"") + + main.stop_connect_container(container) + + container.wait.assert_not_called() + container.stop.assert_called_once_with(timeout=2) + + +def test_stop_connect_container_falls_back_on_errors(): + exec_failure = MagicMock() + exec_failure.exec_run.side_effect = docker.errors.APIError("container is not running") + wait_timeout = MagicMock() + wait_timeout.exec_run.return_value = (0, b"") + wait_timeout.wait.side_effect = requests.exceptions.ReadTimeout("Read timed out.") + + for container in (exec_failure, wait_timeout): + main.stop_connect_container(container) + + container.stop.assert_called_once_with(timeout=2) + + def _decode_segment(segment): return json.loads(base64.urlsafe_b64decode(segment + "=" * (-len(segment) % 4))) @@ -510,6 +589,21 @@ def test_post_bootstrap_raises_on_other_errors(): test_stop_nonexistent_container() print("✓ test_stop_nonexistent_container passed") + test_stop_connect_container_signals_connect() + print("✓ test_stop_connect_container_signals_connect passed") + + test_stop_mode_signals_connect() + print("✓ test_stop_mode_signals_connect passed") + + test_command_mode_signals_connect_after_command() + print("✓ test_command_mode_signals_connect_after_command passed") + + test_stop_connect_container_falls_back_when_connect_not_signaled() + print("✓ test_stop_connect_container_falls_back_when_connect_not_signaled passed") + + test_stop_connect_container_falls_back_on_errors() + print("✓ test_stop_connect_container_falls_back_on_errors passed") + test_make_bootstrap_jwt() print("✓ test_make_bootstrap_jwt passed")