SnapOtter/tests/unit/api/venv-upgrade-stamp.test.ts

import { execFileSync } from "node:child_process";
import { existsSync, mkdirSync, mkdtempSync, readFileSync, rmSync, writeFileSync } from "node:fs";
import { tmpdir } from "node:os";
import { dirname, join, resolve } from "node:path";
import { fileURLToPath } from "node:url";
import { afterEach, beforeEach, describe, expect, it } from "vitest";

/**
 * Unit tests for the venv upgrade-stamp mechanism in docker/entrypoint.sh.
 *
 * The entrypoint compares /opt/venv/.venv-version (baked into the image at
 * build time) against /data/ai/venv/.venv-version (persisted on the volume).
 * When missing or mismatched, the venv is nuked and re-copied so upgraded
 * base packages take effect immediately.
 *
 * These tests exercise the bootstrap logic in isolation using temp dirs.
 */

let root: string;
let optVenv: string;
let dataAi: string;
let aiVenv: string;
let aiVenvTmp: string;
let installedJson: string;

const here = dirname(fileURLToPath(import.meta.url));
const ENTRYPOINT_LIB = resolve(here, "../../../docker/entrypoint-lib.sh");

// Self-contained shell script mirroring the entrypoint bootstrap block.
// Uses env vars for paths so we can point at temp directories.
const BOOTSTRAP_SCRIPT = `
#!/bin/sh
set -e

. "$TEST_ENTRYPOINT_LIB"

AI_VENV="$TEST_AI_VENV"
AI_VENV_TMP="$TEST_AI_VENV_TMP"
OPT_VENV="$TEST_OPT_VENV"

if [ -d "$AI_VENV_TMP" ]; then
  echo "CLEANUP_INTERRUPTED"
  rm -rf "$AI_VENV_TMP"
fi

if [ -d "$OPT_VENV" ]; then
  NEED_BOOTSTRAP=false

  if [ ! -d "$AI_VENV" ]; then
    NEED_BOOTSTRAP=true
    echo "FIRST_RUN"
  elif [ -f "$OPT_VENV/.venv-version" ]; then
    IMAGE_STAMP=$(cat "$OPT_VENV/.venv-version")
    CURRENT_STAMP=""
    if [ -f "$AI_VENV/.venv-version" ]; then
      CURRENT_STAMP=$(cat "$AI_VENV/.venv-version")
    fi
    if [ "$CURRENT_STAMP" != "$IMAGE_STAMP" ]; then
      NEED_BOOTSTRAP=true
      echo "STAMP_MISMATCH"
    fi
  fi

  if [ "$NEED_BOOTSTRAP" = true ]; then
    rm -rf "$AI_VENV"
    cp -r "$OPT_VENV" "$AI_VENV_TMP"
    mv "$AI_VENV_TMP" "$AI_VENV"
    rewrite_venv_paths "$AI_VENV" "$OPT_VENV" "$AI_VENV"
    if [ -f "$TEST_INSTALLED_JSON" ]; then
      echo '{"bundles":{}}' > "$TEST_INSTALLED_JSON"
      echo "BUNDLES_RESET"
    fi
    echo "VENV_READY"
  else
    rewrite_venv_paths "$AI_VENV" "$OPT_VENV" "$AI_VENV"
    echo "SKIP"
  fi
else
  echo "NO_OPT_VENV"
fi
`;

function runBootstrap(): string {
  return execFileSync("/bin/sh", ["-c", BOOTSTRAP_SCRIPT], {
    env: {
      TEST_AI_VENV: aiVenv,
      TEST_AI_VENV_TMP: aiVenvTmp,
      TEST_OPT_VENV: optVenv,
      TEST_INSTALLED_JSON: installedJson,
      TEST_ENTRYPOINT_LIB: ENTRYPOINT_LIB,
      PATH: process.env.PATH,
    },
    encoding: "utf-8",
  }).trim();
}

beforeEach(() => {
  root = mkdtempSync(join(tmpdir(), "venv-stamp-"));
  optVenv = join(root, "opt-venv");
  dataAi = join(root, "data-ai");
  aiVenv = join(dataAi, "venv");
  aiVenvTmp = join(dataAi, "venv.bootstrapping");
  installedJson = join(dataAi, "installed.json");

  // Simulate /opt/venv with a stamp and a marker file
  mkdirSync(optVenv, { recursive: true });
  mkdirSync(join(optVenv, "bin"), { recursive: true });
  writeFileSync(join(optVenv, ".venv-version"), "abc123\n");
  writeFileSync(join(optVenv, "marker.txt"), "base-package-content");
  writeFileSync(join(optVenv, "bin", "pip"), `#!${optVenv}/bin/python3\n`);
  writeFileSync(join(optVenv, "bin", "activate"), `VIRTUAL_ENV=${optVenv}\n`);
  writeFileSync(join(optVenv, "pyvenv.cfg"), `command = python3 -m venv ${optVenv}\n`);

  mkdirSync(dataAi, { recursive: true });
});

afterEach(() => {
  rmSync(root, { recursive: true, force: true });
});

describe("venv upgrade stamp", () => {
  it("bootstraps on first run when no venv exists", () => {
    const output = runBootstrap();
    expect(output).toContain("FIRST_RUN");
    expect(output).toContain("VENV_READY");
    expect(existsSync(join(aiVenv, ".venv-version"))).toBe(true);
    expect(readFileSync(join(aiVenv, "marker.txt"), "utf-8")).toBe("base-package-content");
    expect(readFileSync(join(aiVenv, "bin", "pip"), "utf-8")).toContain(aiVenv);
    expect(readFileSync(join(aiVenv, "bin", "activate"), "utf-8")).toContain(aiVenv);
    expect(readFileSync(join(aiVenv, "pyvenv.cfg"), "utf-8")).toContain(aiVenv);
    expect(readFileSync(join(aiVenv, "bin", "pip"), "utf-8")).not.toContain(optVenv);
  });

  it("skips when stamps match", () => {
    // Simulate an existing venv with matching stamp
    mkdirSync(aiVenv, { recursive: true });
    writeFileSync(join(aiVenv, ".venv-version"), "abc123\n");

    const output = runBootstrap();
    expect(output).toBe("SKIP");
  });

  it("repairs copied venv paths even when stamps already match", () => {
    mkdirSync(join(aiVenv, "bin"), { recursive: true });
    writeFileSync(join(aiVenv, ".venv-version"), "abc123\n");
    writeFileSync(join(aiVenv, "bin", "pip"), `#!${optVenv}/bin/python3\n`);
    writeFileSync(join(aiVenv, "bin", "activate"), `VIRTUAL_ENV=${optVenv}\n`);
    writeFileSync(join(aiVenv, "pyvenv.cfg"), `command = python3 -m venv ${optVenv}\n`);

    const output = runBootstrap();
    expect(output).toBe("SKIP");
    for (const file of [
      join(aiVenv, "bin", "pip"),
      join(aiVenv, "bin", "activate"),
      join(aiVenv, "pyvenv.cfg"),
    ]) {
      const content = readFileSync(file, "utf-8");
      expect(content).toContain(aiVenv);
      expect(content).not.toContain(optVenv);
    }
  });

  it("refreshes when stamps differ (upgrade)", () => {
    // Simulate stale venv with old stamp
    mkdirSync(aiVenv, { recursive: true });
    writeFileSync(join(aiVenv, ".venv-version"), "old-hash\n");
    writeFileSync(join(aiVenv, "stale-file.txt"), "should-be-removed");

    const output = runBootstrap();
    expect(output).toContain("STAMP_MISMATCH");
    expect(output).toContain("VENV_READY");
    // Old content replaced with fresh copy
    expect(existsSync(join(aiVenv, "stale-file.txt"))).toBe(false);
    expect(readFileSync(join(aiVenv, "marker.txt"), "utf-8")).toBe("base-package-content");
    expect(readFileSync(join(aiVenv, ".venv-version"), "utf-8")).toBe("abc123\n");
  });

  it("refreshes when venv has no stamp (pre-stamp image upgrade)", () => {
    // Simulate old venv without any stamp file
    mkdirSync(aiVenv, { recursive: true });
    writeFileSync(join(aiVenv, "old-pkg.txt"), "legacy");

    const output = runBootstrap();
    expect(output).toContain("STAMP_MISMATCH");
    expect(output).toContain("VENV_READY");
    expect(existsSync(join(aiVenv, "old-pkg.txt"))).toBe(false);
  });

  it("resets installed.json when refreshing", () => {
    // Simulate existing venv + installed bundles
    mkdirSync(aiVenv, { recursive: true });
    writeFileSync(join(aiVenv, ".venv-version"), "old-hash\n");
    writeFileSync(
      installedJson,
      JSON.stringify({
        bundles: {
          "background-removal": {
            version: "1.0.0",
            installedAt: "2025-01-01T00:00:00Z",
          },
        },
      }),
    );

    const output = runBootstrap();
    expect(output).toContain("BUNDLES_RESET");
    const data = JSON.parse(readFileSync(installedJson, "utf-8"));
    expect(data).toEqual({ bundles: {} });
  });

  it("does not reset installed.json when no bundles were installed", () => {
    // No installed.json exists
    mkdirSync(aiVenv, { recursive: true });
    writeFileSync(join(aiVenv, ".venv-version"), "old-hash\n");

    const output = runBootstrap();
    expect(output).toContain("VENV_READY");
    expect(output).not.toContain("BUNDLES_RESET");
  });

  it("does nothing when /opt/venv has no stamp (old image)", () => {
    // Remove the stamp from the image venv
    rmSync(join(optVenv, ".venv-version"));
    mkdirSync(aiVenv, { recursive: true });
    writeFileSync(join(aiVenv, "existing.txt"), "keep");

    const output = runBootstrap();
    expect(output).toBe("SKIP");
    // Existing venv untouched
    expect(existsSync(join(aiVenv, "existing.txt"))).toBe(true);
  });

  it("cleans up interrupted bootstrap from previous start", () => {
    mkdirSync(aiVenvTmp, { recursive: true });
    writeFileSync(join(aiVenvTmp, "partial.txt"), "incomplete");

    const output = runBootstrap();
    expect(output).toContain("CLEANUP_INTERRUPTED");
    expect(existsSync(aiVenvTmp)).toBe(false);
  });

  it("does nothing when /opt/venv does not exist", () => {
    rmSync(optVenv, { recursive: true });
    const output = runBootstrap();
    expect(output).toBe("NO_OPT_VENV");
  });
});