appfactory-worker/tests/test_execute_job_streaming.py

"""Regression tests for job output streaming in execute_job().

The subprocess writes into OS pipes with a limited buffer. If the worker does
not read stdout and stderr while the process runs, the child blocks on write
and the job never finishes. These tests run a child that produces far more
output than any pipe buffer can hold.
"""

import json
import os
import sqlite3
import sys
import tempfile
import threading
import unittest
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from app import main  # noqa: E402

# Much larger than the pipe buffer on Linux (64 KiB) and Windows (4 KiB).
CHUNK_LINES = 4000
LINE = "x" * 256

CHILD_SCRIPT = r"""
import sys

line = "x" * 256

for i in range({lines}):
    sys.stdout.write("out %d %s\n" % (i, line))
    sys.stderr.write("err %d %s\n" % (i, line))

sys.stdout.flush()
sys.stderr.flush()
sys.exit({exit_code})
"""

SCHEMA = """
CREATE TABLE jobs (
    id INTEGER PRIMARY KEY,
    type TEXT NOT NULL,
    target_type TEXT,
    target_id TEXT,
    status TEXT NOT NULL,
    payload_json TEXT,
    source TEXT,
    created_by_user_id INTEGER,
    created_by_username TEXT,
    created_by_display_name TEXT,
    worker_id TEXT,
    started_at TEXT,
    finished_at TEXT,
    error_text TEXT,
    result_json TEXT
);

CREATE TABLE job_logs (
    id INTEGER PRIMARY KEY,
    job_id INTEGER NOT NULL,
    stream TEXT NOT NULL,
    message TEXT,
    created_at TEXT
);

CREATE TABLE deployment_locks (
    target_type TEXT NOT NULL,
    target_id TEXT NOT NULL,
    job_id INTEGER,
    worker_id TEXT,
    acquired_at TEXT,
    expires_at TEXT,
    PRIMARY KEY (target_type, target_id)
);

CREATE TABLE workers (
    id TEXT PRIMARY KEY,
    status TEXT,
    started_at TEXT,
    last_seen_at TEXT,
    current_job_id INTEGER,
    metadata_json TEXT
);

CREATE TABLE alert_events (
    id INTEGER PRIMARY KEY,
    status TEXT,
    error_text TEXT,
    processed_at TEXT
);
"""


class ExecuteJobStreamingTest(unittest.TestCase):
    def setUp(self):
        self._tmp_dir = tempfile.TemporaryDirectory()
        self.addCleanup(self._tmp_dir.cleanup)

        db_file = Path(self._tmp_dir.name) / "appfactory.db"

        con = sqlite3.connect(db_file)
        con.executescript(SCHEMA)
        con.commit()
        con.close()

        self._original_db_file = main.DB_FILE
        self._original_command_for_job = main.command_for_job
        self._original_cancel_check = main.CANCEL_CHECK_SECONDS

        main.DB_FILE = db_file
        main.CANCEL_CHECK_SECONDS = 1

        self.addCleanup(self._restore)

        self.db_file = db_file

    def _restore(self):
        main.DB_FILE = self._original_db_file
        main.command_for_job = self._original_command_for_job
        main.CANCEL_CHECK_SECONDS = self._original_cancel_check

    def _connect(self):
        con = sqlite3.connect(self.db_file)
        con.row_factory = sqlite3.Row
        return con

    def _insert_job(self, job_type: str, target_id: str) -> dict:
        con = self._connect()
        cur = con.execute(
            """
            INSERT INTO jobs (type, target_type, target_id, status, payload_json, source)
            VALUES (?, ?, ?, 'running', '{}', 'test')
            """,
            (job_type, job_type, target_id),
        )
        job_id = cur.lastrowid
        con.commit()

        row = con.execute("SELECT * FROM jobs WHERE id = ?", (job_id,)).fetchone()
        con.close()

        return dict(row)

    def _use_child(self, exit_code: int):
        script = CHILD_SCRIPT.format(lines=CHUNK_LINES, exit_code=exit_code)
        script_file = Path(self._tmp_dir.name) / f"child_{exit_code}.py"
        script_file.write_text(script, encoding="utf-8")

        main.command_for_job = lambda job: [sys.executable, str(script_file)]

    def _run_job(self, job: dict, timeout: float = 120.0):
        error = {}

        def runner():
            try:
                main.execute_job(job)
            except BaseException as exc:  # pragma: no cover - reported below
                error["exc"] = exc

        thread = threading.Thread(target=runner, daemon=True)
        thread.start()
        thread.join(timeout)

        self.assertFalse(
            thread.is_alive(),
            f"execute_job did not finish within {timeout} s (deadlock on a full pipe buffer)",
        )

        if "exc" in error:
            raise error["exc"]

    def _job_row(self, job_id: int) -> sqlite3.Row:
        con = self._connect()
        row = con.execute("SELECT * FROM jobs WHERE id = ?", (job_id,)).fetchone()
        con.close()
        return row

    def _log_text(self, job_id: int, stream: str) -> str:
        con = self._connect()
        rows = con.execute(
            "SELECT message FROM job_logs WHERE job_id = ? AND stream = ? ORDER BY id ASC",
            (job_id, stream),
        ).fetchall()
        con.close()
        return "".join(row["message"] for row in rows)

    def test_large_output_on_both_streams_finishes(self):
        """stdout and stderr are drained while the process runs, so it can exit."""
        self._use_child(exit_code=0)

        job = self._insert_job("deploy_app", "demo-app")
        self._run_job(job)

        row = self._job_row(job["id"])
        self.assertEqual(row["status"], "success")

        result = json.loads(row["result_json"])
        self.assertEqual(result["returncode"], 0)

        expected_length = CHUNK_LINES * (len(LINE) + len("out 0 \n"))
        self.assertGreater(result["stdout_length"], 64 * 1024)
        self.assertGreater(result["stderr_length"], 64 * 1024)

        stdout_text = self._log_text(job["id"], "stdout")
        stderr_text = self._log_text(job["id"], "stderr")

        self.assertEqual(len(stdout_text), result["stdout_length"])
        self.assertEqual(len(stderr_text), result["stderr_length"])
        self.assertGreaterEqual(len(stdout_text), expected_length - CHUNK_LINES * 4)

        # Nothing may be lost and nothing stored twice.
        self.assertEqual(stdout_text.count("\nout "), CHUNK_LINES - 1)
        self.assertEqual(stderr_text.count("\nerr "), CHUNK_LINES - 1)

        self.assertTrue(stdout_text.startswith("out 0 "))
        self.assertIn(f"out {CHUNK_LINES - 1} ", stdout_text)
        self.assertIn(f"err {CHUNK_LINES - 1} ", stderr_text)

    def test_large_output_with_failure_is_marked_failed(self):
        """The return code survives the streaming rewrite."""
        self._use_child(exit_code=3)

        job = self._insert_job("run_script", "cleanup.sh")
        self._run_job(job)

        row = self._job_row(job["id"])
        self.assertEqual(row["status"], "failed")
        self.assertIn("3", row["error_text"])

        result = json.loads(row["result_json"])
        self.assertEqual(result["returncode"], 3)

        self.assertEqual(
            self._log_text(job["id"], "stdout").count("\nout "),
            CHUNK_LINES - 1,
        )
        self.assertEqual(
            self._log_text(job["id"], "stderr").count("\nerr "),
            CHUNK_LINES - 1,
        )


if __name__ == "__main__":
    unittest.main(verbosity=2)