From af717d691a76852930a60b10a448dcb2b339e999 Mon Sep 17 00:00:00 2001 From: Ricardo Boni Date: Tue, 25 Aug 2026 23:12:28 -0400 Subject: [PATCH] fix(installer): give TestGen time to stop before docker kills it MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit TestGen's scheduler stops cooperatively: on SIGTERM it asks running jobs to stop at a checkpoint, waits up to TG_JOB_SHUTDOWN_TIMEOUT (default 60s), and records where each one got to so the next start continues from there. Docker's default stop grace period is 10s, so a stop, restart, or upgrade killed the container mid-wait. The checkpoint was never written, the run was canceled, and the next launch started it over from nothing. Measured against a profiling run whose columns take 30s: with 10s the container is SIGKILLed and the next start cancels the run, discarding 38 columns of work; with 90s it stops cleanly in 53s and resumes. 90s clears the 60s wait plus the final write, and matches what TestGen ships in its own compose files and Helm chart. Applied in three places, since the compose template alone reaches only new installs: - The generated compose file, for fresh Docker installs. - tg upgrade, which patches the compose file in place — existing installs keep their file forever, so this is the only path that reaches them. Scoped to the engine service by indentation, and matched on the image key rather than its value, so a private mirror is patched too. - Pip mode, where stop_app_tree force-killed the tree after 10s. It now waits the same grace period, reports whether the stop was clean, and swallows a second Ctrl+C rather than letting it re-signal the tree — TestGen reads a second signal as "hurry up" and cuts the job mid-pause. The force-kill path also sweeps orphans now: run-app all starts its ui and scheduler children in their own sessions, so killpg on the parent left them holding the port and the data directory, breaking the next tg start. Windows still force-kills, as it has no catchable signal to deliver without CREATE_NEW_PROCESS_GROUP. The Ctrl+C messaging is gated on that so it neither promises a wait it cannot honor nor warns about a lost job on every stop. Co-Authored-By: Claude Opus 5 (1M context) --- dk-installer.py | 150 ++++++++++++++++++++++++++---- tests/conftest.py | 11 +++ tests/test_tg_install.py | 19 ++++ tests/test_tg_start.py | 195 +++++++++++++++++++++++++++++++++++++-- tests/test_tg_upgrade.py | 171 +++++++++++++++++++++++++++++++++- 5 files changed, 517 insertions(+), 29 deletions(-) diff --git a/dk-installer.py b/dk-installer.py index 02f2ba5..01b602d 100755 --- a/dk-installer.py +++ b/dk-installer.py @@ -67,6 +67,10 @@ TESTGEN_LOG_FILE_PATH = pathlib.Path.home() / ".testgen" / "logs" / "app.log" TESTGEN_CONFIG_ENV_PATH = pathlib.Path.home() / ".testgen" / "config.env" TESTGEN_APP_READY_TIMEOUT = 120 +# Seconds TestGen is given to stop. Must exceed TG_JOB_SHUTDOWN_TIMEOUT (default 60s) plus +# the time the process needs to record what stopped, or the scheduler is killed mid-wait and +# a running job is cut instead of stopping at a checkpoint. +TESTGEN_STOP_GRACE_PERIOD = 90 INSTALL_MARKER_FILE = "dk-{}-install.json" INSTALL_MODE_DOCKER = "docker" INSTALL_MODE_PIP = "pip" @@ -2079,6 +2083,32 @@ def execute(self, args): CONSOLE.msg("Observability Heartbeat demo stopped.") +def find_in_block(contents: str, block: str, key: str) -> typing.Optional[re.Match]: + """Find the ``key:`` line inside the compose ``block:`` mapping, or None. + + Scans by indentation rather than parsing YAML — enough for the block-style files the + installer writes, and it avoids a runtime dependency. Offsets on the returned match + are absolute, so callers can splice around it; group 1 is the key's indent. + + Scoping to a block is the point: the same key can appear on several services, and + only ``engine`` runs the scheduler. Deliberately says nothing about *which* image a + service uses — ``tg install --image`` accepts any registry. + """ + headers = list(re.finditer(rf"^([ \t]*){re.escape(block)}:[ \t]*$", contents, flags=re.M)) + if not headers: + return None + # Shallowest wins: a name like ``postgres`` is both a service and a nested key under + # another service's ``depends_on``, and it's the service the caller means. + header = min(headers, key=lambda match: len(match.group(1))) + # The block body ends at the first line indented no deeper than the block key itself. + end = len(contents) + for line in re.finditer(r"^([ \t]*)\S.*$", contents[header.end() :], flags=re.M): + if len(line.group(1)) <= len(header.group(1)): + end = header.end() + line.start() + break + return re.compile(rf"^([ \t]+){re.escape(key)}:.*$", flags=re.M).search(contents, header.end(), end) + + class UpdateComposeFileStep(Step): label = "Updating the Docker compose file" @@ -2088,6 +2118,7 @@ def __init__(self): self.update_token = False self.update_base_url = False self.update_api_port = False + self.update_stop_grace = False super().__init__() def pre_execute(self, action, args): @@ -2149,6 +2180,17 @@ def pre_execute(self, action, args): and not re.search(rf"- \d+:{TESTGEN_DEFAULT_API_PORT}\b", contents) ) + # Compose files written before the grace period was added stop the engine after + # docker's 10s default, cutting a running job instead of letting it checkpoint. + # Only count it as a pending change if we can actually place it, or the step + # would report success having silently rewritten the file unchanged. + engine_image = find_in_block(contents, "engine", "image") + if engine_image is None: + LOG.info("No image line in the compose 'engine' service; leaving stop_grace_period alone") + self.update_stop_grace = ( + engine_image is not None and find_in_block(contents, "engine", "stop_grace_period") is None + ) + if not any( ( self.update_version, @@ -2156,6 +2198,7 @@ def pre_execute(self, action, args): self.update_token, self.update_base_url, self.update_api_port, + self.update_stop_grace, ) ): CONSOLE.msg("No changes will be applied.") @@ -2169,6 +2212,7 @@ def execute(self, action, args): self.update_token, self.update_base_url, self.update_api_port, + self.update_stop_grace, ) ): raise SkipStep @@ -2210,6 +2254,15 @@ def execute(self, action, args): new_mapping = f"\n{match.group(1)}- {TESTGEN_DEFAULT_API_PORT}:{TESTGEN_DEFAULT_API_PORT}" contents = contents[0 : match.end()] + new_mapping + contents[match.end() :] + if self.update_stop_grace and (image := find_in_block(contents, "engine", "image")): + indent = image.group(1) + grace = ( + f"\n{indent}# Must exceed TG_JOB_SHUTDOWN_TIMEOUT (default 60s), or docker kills the scheduler" + f"\n{indent}# mid-wait and running jobs are cut instead of stopping at a checkpoint." + f"\n{indent}stop_grace_period: {TESTGEN_STOP_GRACE_PERIOD}s" + ) + contents = contents[: image.end()] + grace + contents[image.end() :] + action.get_compose_file_path(args).write_text(contents) @@ -2309,6 +2362,9 @@ def get_compose_file_contents(self, action, args): services: engine: image: {args.image} + # Must exceed TG_JOB_SHUTDOWN_TIMEOUT (default 60s), or docker kills the scheduler + # mid-wait and running jobs are cut instead of stopping at a checkpoint. + stop_grace_period: {TESTGEN_STOP_GRACE_PERIOD}s container_name: testgen environment: *common-variables volumes: @@ -2524,17 +2580,25 @@ def read_testgen_config_env() -> dict[str, str]: return config -def stop_app_tree(proc: subprocess.Popen, timeout: int = 10) -> None: - """Terminate ``proc`` and all of its descendants. +def supports_graceful_stop() -> bool: + """Whether we can ask the app to stop rather than kill it outright. - Plain ``proc.terminate()`` only kills the parent — pixeltable-pgserver - spawns ``postgres`` children that get orphaned otherwise. Cross-platform: - on Windows we shell out to ``taskkill /F /T``; on POSIX we send SIGTERM - to the whole process group (the parent was started with - ``start_new_session=True``). + POSIX only. On Windows the app is started without ``CREATE_NEW_PROCESS_GROUP``, + so there is no catchable signal we can deliver to it and ``taskkill /F`` is the + only reliable stop — a running job is cut there whatever grace period we would + nominally allow. Callers use this to avoid promising a wait that can't happen. + """ + return platform.system() != "Windows" + + +def force_kill_app_tree(proc: subprocess.Popen, timeout: int = 5) -> None: + """Kill ``proc`` and every descendant, including those outside its process group. + + ``testgen run-app all`` starts its ui/scheduler children in their own sessions, so + they sit outside ``proc``'s process group and outlive a ``killpg`` — killing only + the parent leaves the UI holding its port and postgres holding the data directory, + which then breaks the next ``tg start``. Hence the orphan sweep to finish the job. """ - if proc.poll() is not None: - return if platform.system() == "Windows": with contextlib.suppress(Exception): subprocess.run( @@ -2546,16 +2610,51 @@ def stop_app_tree(proc: subprocess.Popen, timeout: int = 10) -> None: ) else: with contextlib.suppress(Exception): - os.killpg(os.getpgid(proc.pid), signal.SIGTERM) + os.killpg(os.getpgid(proc.pid), signal.SIGKILL) + # Backstop in case the platform kill above didn't land (e.g. taskkill denied). + proc.kill() + with contextlib.suppress(subprocess.TimeoutExpired): + proc.wait(timeout=timeout) + stop_standalone_orphans() + + +def stop_app_tree(proc: subprocess.Popen, timeout: int = 10) -> bool: + """Terminate ``proc`` and all of its descendants. Returns whether it stopped on its own. + + Plain ``proc.terminate()`` only kills the parent — pixeltable-pgserver + spawns ``postgres`` children that get orphaned otherwise. Cross-platform: + on Windows we shell out to ``taskkill /F /T``; on POSIX we send SIGTERM + to the whole process group (the parent was started with + ``start_new_session=True``) and let it forward the signal to its children. + + ``timeout`` is how long the tree is given to shut down cooperatively before + it is force-killed — the caller passes ``TESTGEN_STOP_GRACE_PERIOD`` when a + running job may need to reach a checkpoint first. A second Ctrl+C during + that wait is taken as "stop now" and skips straight to the force-kill. + + Returns ``True`` when the tree exited within ``timeout`` (or was already + gone), ``False`` when it had to be force-killed. See ``supports_graceful_stop`` + for why Windows always reports ``False`` when there was a live process. + """ + if proc.poll() is not None: + return True + + if not supports_graceful_stop(): + force_kill_app_tree(proc, timeout=timeout) + return False + + with contextlib.suppress(Exception): + os.killpg(os.getpgid(proc.pid), signal.SIGTERM) try: proc.wait(timeout=timeout) - except subprocess.TimeoutExpired: - if platform.system() != "Windows": - with contextlib.suppress(Exception): - os.killpg(os.getpgid(proc.pid), signal.SIGKILL) - proc.kill() - with contextlib.suppress(subprocess.TimeoutExpired): - proc.wait(timeout=5) + except (subprocess.TimeoutExpired, KeyboardInterrupt): + # KeyboardInterrupt here is a second Ctrl+C while we were waiting. Swallow it + # rather than letting it unwind into the caller's ``finally``, which would send + # the tree another SIGTERM — TestGen reads a second signal as "hurry up" and + # force-kills the job mid-pause, losing exactly the checkpoint we were waiting for. + force_kill_app_tree(proc) + return False + return True def stop_standalone_orphans() -> None: @@ -2687,8 +2786,21 @@ def start_testgen_app(action, args) -> None: # Reset the cursor to column 0 — the terminal echoed `^C` mid-line. print("") CONSOLE.msg("Stopping TestGen...") - stop_app_tree(proc, timeout=10) - CONSOLE.msg("TestGen stopped.") + graceful = supports_graceful_stop() + if graceful: + # A running profiling job stops at its next checkpoint rather than being cut, + # but that can take up to TG_JOB_SHUTDOWN_TIMEOUT — say so, or the wait reads + # as a hang and the user reaches for a second Ctrl+C. + CONSOLE.msg( + f"Waiting up to {TESTGEN_STOP_GRACE_PERIOD} seconds for running jobs to reach a checkpoint..." + ) + stopped_cleanly = stop_app_tree(proc, timeout=TESTGEN_STOP_GRACE_PERIOD) + # Only worth flagging where we actually offered a grace period. Windows always + # force-kills, so the warning would fire on every stop and mean nothing. + if graceful and not stopped_cleanly: + CONSOLE.msg("TestGen stopped. A job that was still running will restart from the beginning.") + else: + CONSOLE.msg("TestGen stopped.") CONSOLE.msg(f"To start it again, {command_hint(args.prod, 'start', 'Start TestGen')}.") finally: stop_app_tree(proc, timeout=5) diff --git a/tests/conftest.py b/tests/conftest.py index 3d57a0d..41b737a 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -33,6 +33,17 @@ def _no_real_process_group_signals(): yield killpg_mock +@pytest.fixture(autouse=True) +def _no_real_orphan_sweep(): + """``force_kill_app_tree`` finishes with ``stop_standalone_orphans``, which shells out + to a real ``pkill -9 -f 'testgen.*run-app'``. Unpatched, a test exercising the + force-kill path would kill the developer's own running TestGen. Tests that assert on + the sweep override this inside their own ``with patch(...)``. + """ + with patch("tests.installer.stop_standalone_orphans") as mock: + yield mock + + @pytest.fixture def stdout_mock(): return Mock(return_value=[]) diff --git a/tests/test_tg_install.py b/tests/test_tg_install.py index ce2cf5f..85228d8 100644 --- a/tests/test_tg_install.py +++ b/tests/test_tg_install.py @@ -10,6 +10,7 @@ AbortAction, TestGenCreateDockerComposeFileStep, ComposeVerifyExistingInstallStep, + TESTGEN_STOP_GRACE_PERIOD, ) @@ -114,6 +115,24 @@ def test_tg_compose_base_url_custom_port(tg_install_action, start_cmd_mock, stdo assert "TG_UI_BASE_URL: http://localhost:9000" in contents +@pytest.mark.integration +def test_tg_compose_sets_engine_stop_grace_period(tg_install_action, start_cmd_mock, stdout_mock, compose_path): + """Docker's 10s default kills the scheduler mid-shutdown, cutting a running job + instead of letting it stop at a checkpoint.""" + tg_install_action.execute() + compose_content = compose_path.read_text() + + # Only the engine needs it — postgres shuts down on its own quickly. + assert compose_content.count("stop_grace_period") == 1 + lines = compose_content.splitlines() + image_idx = next(i for i, line in enumerate(lines) if "image: datakitchen/dataops-testgen" in line) + grace_idx = next(i for i, line in enumerate(lines) if "stop_grace_period" in line) + # Inside the engine service, right under its image: two comment lines, then the key. + assert grace_idx == image_idx + 3 + indent = lines[image_idx][: -len(lines[image_idx].lstrip())] + assert lines[grace_idx] == f"{indent}stop_grace_period: {TESTGEN_STOP_GRACE_PERIOD}s" + + @pytest.mark.integration def test_tg_compose_base_url_ssl(tg_install_action, start_cmd_mock, stdout_mock, args_mock, compose_path): args_mock.ssl_cert_file = "/path/to/cert.crt" diff --git a/tests/test_tg_start.py b/tests/test_tg_start.py index ecd86ba..a798764 100644 --- a/tests/test_tg_start.py +++ b/tests/test_tg_start.py @@ -1,3 +1,4 @@ +import signal from pathlib import Path from unittest.mock import MagicMock, patch @@ -11,7 +12,9 @@ TestgenStartAction, start_testgen_app, stop_app_tree, + force_kill_app_tree, InstallMarker, + TESTGEN_STOP_GRACE_PERIOD, ) @@ -117,18 +120,66 @@ def test_start_testgen_app_handles_keyboard_interrupt(app_action, args_mock, con patch("tests.installer.resolve_testgen_path", return_value="/bin/testgen"), patch("tests.installer.subprocess.Popen", return_value=proc), patch("tests.installer.wait_for_tcp_port", return_value=True), - patch("tests.installer.stop_app_tree") as stop_mock, + patch("tests.installer.stop_app_tree", return_value=True) as stop_mock, ): start_testgen_app(app_action, args_mock) - # Called once for the keyboard-interrupt branch (timeout=10) and again in - # the ``finally`` cleanup (timeout=5; no-op since proc already stopped). + # Called once for the keyboard-interrupt branch and again in the ``finally`` + # cleanup (timeout=5; no-op since proc already stopped). The first wait gets the + # full grace period so a running job can reach a checkpoint before being killed. assert stop_mock.call_args_list[0].args[0] is proc - assert stop_mock.call_args_list[0].kwargs == {"timeout": 10} - console_msg_mock.assert_any_msg_contains("TestGen stopped") + assert stop_mock.call_args_list[0].kwargs == {"timeout": TESTGEN_STOP_GRACE_PERIOD} + console_msg_mock.assert_any_msg_contains("reach a checkpoint") + console_msg_mock.assert_any_msg_contains("TestGen stopped.") console_msg_mock.assert_any_msg_contains("tg start") +@pytest.mark.unit +def test_start_testgen_app_warns_when_grace_period_expires(app_action, args_mock, console_msg_mock, empty_tg_config): + """A tree that had to be force-killed lost its checkpoint — say so rather than + reporting a clean stop the user can't trust.""" + args_mock.prod = "tg" + + proc = MagicMock() + proc.poll.return_value = None + proc.wait.side_effect = [KeyboardInterrupt(), 0] + + with ( + patch("tests.installer.resolve_testgen_path", return_value="/bin/testgen"), + patch("tests.installer.subprocess.Popen", return_value=proc), + patch("tests.installer.wait_for_tcp_port", return_value=True), + patch("tests.installer.stop_app_tree", return_value=False), + ): + start_testgen_app(app_action, args_mock) + + console_msg_mock.assert_any_msg_contains("will restart from the beginning") + + +@pytest.mark.unit +def test_start_testgen_app_makes_no_grace_promise_on_windows(app_action, args_mock, console_msg_mock, empty_tg_config): + """Windows always force-kills, so promising a wait would be a lie and the + force-kill warning would fire on every single stop, saying nothing.""" + args_mock.prod = "tg" + + proc = MagicMock() + proc.poll.return_value = None + proc.wait.side_effect = [KeyboardInterrupt(), 0] + + with ( + patch("tests.installer.resolve_testgen_path", return_value="/bin/testgen"), + patch("tests.installer.subprocess.Popen", return_value=proc), + patch("tests.installer.wait_for_tcp_port", return_value=True), + patch("tests.installer.supports_graceful_stop", return_value=False), + patch("tests.installer.stop_app_tree", return_value=False), + ): + start_testgen_app(app_action, args_mock) + + printed = " ".join(str(c) for c in console_msg_mock.call_args_list) + assert "reach a checkpoint" not in printed + assert "will restart from the beginning" not in printed + console_msg_mock.assert_any_msg_contains("TestGen stopped.") + + # --- stop_app_tree ------------------------------------------------------------ @@ -197,7 +248,139 @@ def test_stop_app_tree_falls_through_to_kill_on_timeout(): patch("tests.installer.os.killpg"), patch("tests.installer.os.getpgid", return_value=4242), ): - stop_app_tree(proc, timeout=1) + forced = stop_app_tree(proc, timeout=1) + + proc.kill.assert_called_once() + assert forced is False # had to be force-killed + + +@pytest.mark.unit +def test_stop_app_tree_reports_graceful_stop(): + """Exiting within the grace period is the signal the caller uses to decide + whether a running job got to checkpoint.""" + proc = MagicMock() + proc.poll.return_value = None + proc.pid = 4242 + proc.wait.return_value = 0 + + with ( + patch("tests.installer.platform.system", return_value="Linux"), + patch("tests.installer.os.killpg"), + patch("tests.installer.os.getpgid", return_value=4242), + ): + assert stop_app_tree(proc, timeout=3) is True + + +@pytest.mark.unit +def test_stop_app_tree_swallows_second_interrupt(): + """A second Ctrl+C while we wait means 'stop now'. It must force-kill here rather + than escaping to the caller's ``finally``, which would re-signal the tree — TestGen + reads a second signal as 'hurry up' and cuts the job mid-pause.""" + proc = MagicMock() + proc.poll.return_value = None + proc.pid = 4242 + proc.wait.side_effect = [KeyboardInterrupt(), 0] + + with ( + patch("tests.installer.platform.system", return_value="Linux"), + patch("tests.installer.os.killpg") as killpg_mock, + patch("tests.installer.os.getpgid", return_value=4242), + ): + forced = stop_app_tree(proc, timeout=90) + + assert forced is False + proc.kill.assert_called_once() + # SIGTERM first, then the escalation to SIGKILL. + assert [c.args[1] for c in killpg_mock.call_args_list] == [signal.SIGTERM, signal.SIGKILL] + + +@pytest.mark.unit +def test_stop_app_tree_windows_stop_is_always_forced(): + """``taskkill /F`` gives the tree no chance to checkpoint, so the caller must not be + told the stop was graceful.""" + proc = MagicMock() + proc.poll.return_value = None + proc.pid = 4242 + proc.wait.return_value = 0 + + with ( + patch("tests.installer.platform.system", return_value="Windows"), + patch("tests.installer.subprocess.run"), + ): + assert stop_app_tree(proc, timeout=90) is False + + +@pytest.mark.unit +def test_force_kill_sweeps_orphans_outside_the_process_group(): + """``run-app all`` starts its ui/scheduler children in their own sessions, so killpg + on the parent leaves them holding the port and the data dir — breaking the next + ``tg start``. The sweep is what actually finishes them off.""" + proc = MagicMock() + proc.poll.return_value = None + proc.pid = 4242 + + with ( + patch("tests.installer.platform.system", return_value="Linux"), + patch("tests.installer.os.killpg"), + patch("tests.installer.os.getpgid", return_value=4242), + patch("tests.installer.stop_standalone_orphans") as sweep_mock, + ): + force_kill_app_tree(proc) + + sweep_mock.assert_called_once() + + +@pytest.mark.unit +def test_second_interrupt_still_sweeps_orphans(): + """The second-Ctrl+C path force-kills, so it owes the same cleanup.""" + proc = MagicMock() + proc.poll.return_value = None + proc.pid = 4242 + proc.wait.side_effect = [KeyboardInterrupt(), 0] + + with ( + patch("tests.installer.platform.system", return_value="Linux"), + patch("tests.installer.os.killpg"), + patch("tests.installer.os.getpgid", return_value=4242), + patch("tests.installer.stop_standalone_orphans") as sweep_mock, + ): + assert stop_app_tree(proc, timeout=90) is False + + sweep_mock.assert_called_once() + + +@pytest.mark.unit +def test_graceful_stop_does_not_sweep_orphans(): + """A tree that stopped on its own has nothing left behind — don't reach for pkill.""" + proc = MagicMock() + proc.poll.return_value = None + proc.pid = 4242 + proc.wait.return_value = 0 + + with ( + patch("tests.installer.platform.system", return_value="Linux"), + patch("tests.installer.os.killpg"), + patch("tests.installer.os.getpgid", return_value=4242), + patch("tests.installer.stop_standalone_orphans") as sweep_mock, + ): + assert stop_app_tree(proc, timeout=3) is True + + sweep_mock.assert_not_called() + + +@pytest.mark.unit +def test_force_kill_falls_back_when_taskkill_fails(): + """taskkill can be denied (elevated child). Without the proc.kill() backstop the + installer would report a stop that never happened.""" + proc = MagicMock() + proc.poll.return_value = None + proc.pid = 4242 + + with ( + patch("tests.installer.platform.system", return_value="Windows"), + patch("tests.installer.subprocess.run", side_effect=OSError("Access denied")), + ): + force_kill_app_tree(proc, timeout=3) proc.kill.assert_called_once() diff --git a/tests/test_tg_upgrade.py b/tests/test_tg_upgrade.py index dfa5f31..21d57d6 100644 --- a/tests/test_tg_upgrade.py +++ b/tests/test_tg_upgrade.py @@ -9,7 +9,9 @@ AbortAction, CommandFailed, TESTGEN_MAJOR_VERSION, + TESTGEN_STOP_GRACE_PERIOD, TestgenUpgradeAction, + find_in_block, InstallMarker, ) @@ -45,7 +47,12 @@ def tg_upgrade_stdout_side_effect(stdout_mock): yield side_effect -def get_compose_content(*extra_vars): +def get_compose_content(*extra_vars, stop_grace=False): + """A compose file as an older installer would have written it. + + ``stop_grace`` opts into the engine grace period, i.e. a file already current + in that respect — leave it off to model the installs the upgrade has to patch. + """ template = textwrap.dedent(""" name: testgen @@ -63,10 +70,11 @@ def get_compose_content(*extra_vars): services: engine: image: datakitchen/dataops-testgen:v2.14.5 - + {} """) - return template.format(textwrap.indent("\n".join(extra_vars), " ")) + grace = f" stop_grace_period: {TESTGEN_STOP_GRACE_PERIOD}s\n" if stop_grace else "" + return template.format(textwrap.indent("\n".join(extra_vars), " "), grace) def set_version_check_mock(version_check_mock, latest_version): @@ -137,7 +145,7 @@ def test_tg_upgrade_abort( args_mock.skip_verify = False set_version_check_mock(version_check_mock, "1.0.0") initial_compose_content = get_compose_content( - "TG_INSTANCE_ID: test-instance-id", "TG_UI_BASE_URL: http://localhost:8501" + "TG_INSTANCE_ID: test-instance-id", "TG_UI_BASE_URL: http://localhost:8501", stop_grace=True ) compose_path.write_text(initial_compose_content) @@ -239,3 +247,158 @@ def test_tg_upgrade_preserves_existing_base_url( compose_content = compose_path.read_text() assert "TG_UI_BASE_URL: https://custom.example.com" in compose_content assert compose_content.count("TG_UI_BASE_URL") == 1 + + +@pytest.mark.integration +def test_tg_upgrade_adds_stop_grace_period( + tg_upgrade_action, + compose_path, + start_cmd_mock, + tg_upgrade_stdout_side_effect, + args_mock, + version_check_mock, +): + """Existing installs keep their compose file forever — the upgrade is the only + chance to give them the grace period a running job needs to checkpoint.""" + set_version_check_mock(version_check_mock, "1.0.0") + compose_path.write_text(get_compose_content("TG_INSTANCE_ID: test-instance-id")) + + tg_upgrade_action.execute(args_mock) + + compose_content = compose_path.read_text() + lines = compose_content.splitlines() + image_idx = next(i for i, line in enumerate(lines) if "image: datakitchen/dataops-testgen" in line) + grace_idx = next(i for i, line in enumerate(lines) if "stop_grace_period" in line) + # Inside the engine service, right under its image: two comment lines, then the key. + assert grace_idx == image_idx + 3 + indent = lines[image_idx][: -len(lines[image_idx].lstrip())] + assert lines[grace_idx] == f"{indent}stop_grace_period: {TESTGEN_STOP_GRACE_PERIOD}s" + + +@pytest.mark.integration +def test_tg_upgrade_preserves_existing_stop_grace_period( + tg_upgrade_action, + compose_path, + start_cmd_mock, + tg_upgrade_stdout_side_effect, + args_mock, + version_check_mock, +): + """A user who tuned the value keeps it, and repeated upgrades don't stack duplicates.""" + args_mock.skip_verify = True + set_version_check_mock(version_check_mock, "1.1.0") + compose_path.write_text( + get_compose_content("TG_INSTANCE_ID: test-instance-id").replace( + "image: datakitchen/dataops-testgen:v2.14.5", + "image: datakitchen/dataops-testgen:v2.14.5\n stop_grace_period: 300s", + ) + ) + + tg_upgrade_action.execute(args_mock) + + compose_content = compose_path.read_text() + assert "stop_grace_period: 300s" in compose_content + assert compose_content.count("stop_grace_period") == 1 + + +@pytest.mark.integration +def test_tg_upgrade_adds_stop_grace_period_to_custom_image( + tg_upgrade_action, + compose_path, + start_cmd_mock, + tg_upgrade_stdout_side_effect, + args_mock, + version_check_mock, +): + """``tg install --image`` accepts a private mirror, so the anchor can't assume the + image is a datakitchen one — those installs need the grace period just as much.""" + args_mock.skip_verify = True + set_version_check_mock(version_check_mock, "1.1.0") + compose_path.write_text( + get_compose_content("TG_INSTANCE_ID: test-instance-id").replace( + "datakitchen/dataops-testgen:v2.14.5", "registry.internal.example.com/mirror/testgen:v2.14.5" + ) + ) + + tg_upgrade_action.execute(args_mock) + + compose_content = compose_path.read_text() + lines = compose_content.splitlines() + image_idx = next(i for i, line in enumerate(lines) if "image:" in line) + grace_idx = next(i for i, line in enumerate(lines) if "stop_grace_period" in line) + assert grace_idx == image_idx + 3 + assert lines[grace_idx].strip() == f"stop_grace_period: {TESTGEN_STOP_GRACE_PERIOD}s" + + +@pytest.mark.integration +def test_tg_upgrade_ignores_stop_grace_period_on_another_service( + tg_upgrade_action, + compose_path, + start_cmd_mock, + tg_upgrade_stdout_side_effect, + args_mock, + version_check_mock, +): + """A grace period set on postgres says nothing about the service that runs the + scheduler — the engine must still get its own.""" + args_mock.skip_verify = True + set_version_check_mock(version_check_mock, "1.1.0") + compose_path.write_text( + # The stray comment matters too: a mention anywhere else in the file must not make + # the engine look already-patched. + "# note: stop_grace_period is managed by the installer\n" + + get_compose_content("TG_INSTANCE_ID: test-instance-id") + + "\n postgres:\n image: postgres:14.1-alpine\n stop_grace_period: 30s\n" + ) + + tg_upgrade_action.execute(args_mock) + + compose_content = compose_path.read_text() + engine_block, _, postgres_block = compose_content.partition(" postgres:") + assert f"stop_grace_period: {TESTGEN_STOP_GRACE_PERIOD}s" in engine_block + # The user's postgres value is left exactly as they set it. + assert "stop_grace_period: 30s" in postgres_block + # engine's, postgres', and the stray comment. + assert compose_content.count("stop_grace_period") == 3 + + +COMPOSE_TWO_SERVICES = """name: testgen +# a stray mention of stop_grace_period above the services section +services: + engine: + image: datakitchen/dataops-testgen:v5 + depends_on: + postgres: + condition: service_healthy + + postgres: + image: postgres:14.1-alpine + stop_grace_period: 30s +""" + + +@pytest.mark.unit +@pytest.mark.parametrize( + "block, key, expected", + ( + ("engine", "image", "image: datakitchen/dataops-testgen:v5"), + ("postgres", "image", "image: postgres:14.1-alpine"), + ("postgres", "stop_grace_period", "stop_grace_period: 30s"), + # Scoping is the whole point: postgres' grace period is not the engine's, and a + # mention in a comment above `services:` is not a setting on any service. + ("engine", "stop_grace_period", None), + ("engine", "nonexistent", None), + ("nonexistent", "image", None), + ), +) +def test_find_in_block_is_scoped_to_the_block(block, key, expected): + match = find_in_block(COMPOSE_TWO_SERVICES, block, key) + assert (match.group(0).strip() if match else None) == expected + + +@pytest.mark.unit +def test_find_in_block_offsets_are_absolute(): + """Callers splice around the match, so its offsets must index the whole file.""" + match = find_in_block(COMPOSE_TWO_SERVICES, "postgres", "image") + assert COMPOSE_TWO_SERVICES[match.start() : match.end()] == " image: postgres:14.1-alpine" + assert match.group(1) == " "