@useatlas/create 0.0.6 → 0.0.7

package/templates/docker/src/lib/tools/__tests__/python-nsjail.test.ts

@@ -1,515 +0,0 @@
-import { describe, expect, it, beforeEach, afterEach, mock } from "bun:test";
-
-// ---------------------------------------------------------------------------
-// Mocks
-// ---------------------------------------------------------------------------
-
-// Mock logger to avoid side effects
-mock.module("@atlas/api/lib/logger", () => ({
-  createLogger: () => ({
-    debug: () => {},
-    info: () => {},
-    warn: () => {},
-    error: () => {},
-  }),
-}));
-
-// Mock tracing
-mock.module("@atlas/api/lib/tracing", () => ({
-  withSpan: async (_name: string, _attrs: unknown, fn: () => Promise<unknown>) => fn(),
-}));
-
-// Mock fs operations to avoid real filesystem access
-mock.module("fs", () => ({
-  mkdirSync: () => undefined,
-  writeFileSync: () => undefined,
-  rmSync: () => undefined,
-  accessSync: () => undefined,
-  constants: { X_OK: 1, R_OK: 4 },
-}));
-
-// Track Bun.spawn calls
-let spawnCalls: { args: unknown[]; options: unknown }[] = [];
-let spawnResult: {
-  stdin: { write: (d: string) => void; end: () => void };
-  stdout: ReadableStream;
-  stderr: ReadableStream;
-  exited: Promise<number>;
-  kill: () => void;
-};
-
-function makeStream(text: string): ReadableStream {
-  return new ReadableStream({
-    start(controller) {
-      controller.enqueue(new TextEncoder().encode(text));
-      controller.close();
-    },
-  });
-}
-
-function setSpawnResult(stdout: string, stderr: string, exitCode: number) {
-  spawnResult = {
-    stdin: { write: () => {}, end: () => {} },
-    stdout: makeStream(stdout),
-    stderr: makeStream(stderr),
-    exited: Promise.resolve(exitCode),
-    kill: () => {},
-  };
-}
-
-// Default: successful empty output
-setSpawnResult("", "", 0);
-
-Bun.spawn = ((...args: unknown[]) => {
-  spawnCalls.push({ args: [args[0]], options: args[1] });
-  return spawnResult;
-}) as unknown as typeof Bun.spawn;
-
-const { buildPythonNsjailArgs, createPythonNsjailBackend } = await import("@atlas/api/lib/tools/python-nsjail");
-
-// ---------------------------------------------------------------------------
-// Tests
-// ---------------------------------------------------------------------------
-
-describe("buildPythonNsjailArgs", () => {
-  const originalEnv = { ...process.env };
-
-  afterEach(() => {
-    process.env = { ...originalEnv };
-  });
-
-  it("constructs correct nsjail args with Python-specific config", () => {
-    const args = buildPythonNsjailArgs(
-      "/usr/local/bin/nsjail",
-      "/tmp/pyexec-test",
-      "/tmp/pyexec-test/user_code.py",
-      "/tmp/pyexec-test/wrapper.py",
-      "/tmp/pyexec-test/charts",
-      "__ATLAS_RESULT_test__",
-    );
-
-    // Basic nsjail mode
-    expect(args[0]).toBe("/usr/local/bin/nsjail");
-    expect(args).toContain("--mode");
-    expect(args).toContain("o");
-
-    // Python binary bind-mounts
-    expect(args).toContain("/usr/local/bin");
-    expect(args).toContain("/usr/local/lib");
-
-    // Code and chart bind-mounts
-    const rFlags = args.reduce<string[]>((acc, v, i) => {
-      if (v === "-R" && typeof args[i + 1] === "string") acc.push(args[i + 1] as string);
-      return acc;
-    }, []);
-    expect(rFlags).toContain("/tmp/pyexec-test/wrapper.py:/tmp/wrapper.py");
-    expect(rFlags).toContain("/tmp/pyexec-test/user_code.py:/tmp/user_code.py");
-
-    // Chart dir is bind-mounted writable (-B)
-    const bIndex = args.indexOf("-B");
-    expect(bIndex).toBeGreaterThan(-1);
-    expect(args[bIndex + 1]).toBe("/tmp/pyexec-test/charts:/tmp/charts");
-
-    // Resource limits — higher defaults for Python
-    const memIndex = args.indexOf("--rlimit_as");
-    expect(memIndex).toBeGreaterThan(-1);
-    expect(args[memIndex + 1]).toBe("512");
-
-    const tIndex = args.indexOf("-t");
-    expect(tIndex).toBeGreaterThan(-1);
-    expect(args[tIndex + 1]).toBe("30");
-
-    const nprocIndex = args.indexOf("--rlimit_nproc");
-    expect(nprocIndex).toBeGreaterThan(-1);
-    expect(args[nprocIndex + 1]).toBe("16");
-
-    // File size limit for chart output
-    const fsizeIndex = args.indexOf("--rlimit_fsize");
-    expect(fsizeIndex).toBeGreaterThan(-1);
-    expect(args[fsizeIndex + 1]).toBe("50");
-
-    // Security: run as nobody
-    const uIndex = args.indexOf("-u");
-    expect(uIndex).toBeGreaterThan(-1);
-    expect(args[uIndex + 1]).toBe("65534");
-
-    const gIndex = args.indexOf("-g");
-    expect(gIndex).toBeGreaterThan(-1);
-    expect(args[gIndex + 1]).toBe("65534");
-
-    // stdin passthrough
-    expect(args).toContain("--pass_fd");
-    const passFdIndex = args.indexOf("--pass_fd");
-    expect(args[passFdIndex + 1]).toBe("0");
-
-    // Python execution command
-    const dashDash = args.indexOf("--");
-    expect(args[dashDash + 1]).toBe("/usr/bin/python3");
-    expect(args[dashDash + 2]).toBe("/tmp/wrapper.py");
-    expect(args[dashDash + 3]).toBe("/tmp/user_code.py");
-
-    // Suppress logs
-    expect(args).toContain("--quiet");
-
-    // /proc mount
-    const procIndex = args.indexOf("--proc_path");
-    expect(procIndex).toBeGreaterThan(-1);
-    expect(args[procIndex + 1]).toBe("/proc");
-  });
-
-  it("applies configurable resource limits", () => {
-    process.env.ATLAS_NSJAIL_TIME_LIMIT = "60";
-    process.env.ATLAS_NSJAIL_MEMORY_LIMIT = "1024";
-
-    const args = buildPythonNsjailArgs(
-      "/usr/local/bin/nsjail",
-      "/tmp/test",
-      "/tmp/test/code.py",
-      "/tmp/test/wrapper.py",
-      "/tmp/test/charts",
-      "marker",
-    );
-
-    const tIndex = args.indexOf("-t");
-    expect(args[tIndex + 1]).toBe("60");
-
-    const memIndex = args.indexOf("--rlimit_as");
-    expect(args[memIndex + 1]).toBe("1024");
-  });
-});
-
-describe("createPythonNsjailBackend", () => {
-  const originalEnv = { ...process.env };
-
-  beforeEach(() => {
-    spawnCalls = [];
-    setSpawnResult("", "", 0);
-  });
-
-  afterEach(() => {
-    process.env = { ...originalEnv };
-  });
-
-  it("parses structured result from marker line", async () => {
-    // Capture the marker from the env passed to Bun.spawn, then verify parsing
-    const backend = createPythonNsjailBackend("/usr/local/bin/nsjail");
-
-    // We need the marker from the spawn call, so use a custom spawn mock
-    // that echoes a result with the actual marker
-    const savedSpawn = Bun.spawn;
-    Bun.spawn = ((...args: unknown[]) => {
-      spawnCalls.push({ args: [args[0]], options: args[1] });
-      const opts = args[1] as { env: Record<string, string> };
-      const marker = opts.env.ATLAS_RESULT_MARKER;
-      return {
-        stdin: { write: () => {}, end: () => {} },
-        stdout: makeStream(`${marker}{"success":true,"output":"hello world"}\n`),
-        stderr: makeStream(""),
-        exited: Promise.resolve(0),
-        kill: () => {},
-      };
-    }) as unknown as typeof Bun.spawn;
-
-    const result = await backend.exec('print("hello")');
-
-    expect(result.success).toBe(true);
-    if (result.success) {
-      expect(result.output).toBe("hello world");
-    }
-
-    // Env should have no secrets
-    const spawnOpts = spawnCalls[0].options as { env: Record<string, string> };
-    expect(spawnOpts.env.MPLBACKEND).toBe("Agg");
-    expect(spawnOpts.env.ATLAS_CHART_DIR).toBe("/tmp/charts");
-    expect(spawnOpts.env.ATLAS_RESULT_MARKER).toBeDefined();
-    expect(spawnOpts.env).not.toHaveProperty("ATLAS_DATASOURCE_URL");
-    expect(spawnOpts.env).not.toHaveProperty("ANTHROPIC_API_KEY");
-
-    Bun.spawn = savedSpawn;
-  });
-
-  it("passes data as stdin when provided", async () => {
-    let stdinWritten = "";
-    spawnResult = {
-      stdin: {
-        write: (d: string) => { stdinWritten = d; },
-        end: () => {},
-      },
-      stdout: makeStream(""),
-      stderr: makeStream(""),
-      exited: Promise.resolve(0),
-      kill: () => {},
-    };
-
-    const backend = createPythonNsjailBackend("/usr/local/bin/nsjail");
-    await backend.exec("print(df.head())", { columns: ["a", "b"], rows: [[1, 2]] });
-
-    const parsed = JSON.parse(stdinWritten);
-    expect(parsed.columns).toEqual(["a", "b"]);
-    expect(parsed.rows).toEqual([[1, 2]]);
-  });
-
-  it("returns error when no result marker in output", async () => {
-    setSpawnResult("some random output", "ImportError: no module named foobar", 1);
-
-    const backend = createPythonNsjailBackend("/usr/local/bin/nsjail");
-    const result = await backend.exec("import foobar");
-
-    expect(result.success).toBe(false);
-    if (!result.success) {
-      expect(result.error).toContain("ImportError");
-    }
-  });
-
-  it("returns error when process killed by signal (SIGKILL)", async () => {
-    setSpawnResult("", "", 137); // 128 + 9 = SIGKILL
-
-    const backend = createPythonNsjailBackend("/usr/local/bin/nsjail");
-    const result = await backend.exec("while True: pass");
-
-    expect(result.success).toBe(false);
-    if (!result.success) {
-      expect(result.error).toContain("killed");
-    }
-  });
-
-  it("returns signal-specific error for non-SIGKILL signals", async () => {
-    setSpawnResult("", "Segmentation fault", 139); // 128 + 11 = SIGSEGV
-
-    const backend = createPythonNsjailBackend("/usr/local/bin/nsjail");
-    const result = await backend.exec("bad code");
-
-    expect(result.success).toBe(false);
-    if (!result.success) {
-      expect(result.error).toContain("SIGSEGV");
-      expect(result.error).toContain("Segmentation fault");
-    }
-  });
-
-  it("returns error when nsjail spawn fails", async () => {
-    const savedSpawn = Bun.spawn;
-    Bun.spawn = (() => {
-      throw new Error("spawn failed: permission denied");
-    }) as unknown as typeof Bun.spawn;
-
-    const backend = createPythonNsjailBackend("/usr/local/bin/nsjail");
-    const result = await backend.exec("print(1)");
-
-    expect(result.success).toBe(false);
-    if (!result.success) {
-      expect(result.error).toContain("nsjail infrastructure error");
-    }
-
-    Bun.spawn = savedSpawn;
-  });
-
-  it("sends empty string on stdin when no data", async () => {
-    let stdinWritten = "";
-    spawnResult = {
-      stdin: {
-        write: (d: string) => { stdinWritten = d; },
-        end: () => {},
-      },
-      stdout: makeStream(""),
-      stderr: makeStream(""),
-      exited: Promise.resolve(0),
-      kill: () => {},
-    };
-
-    const backend = createPythonNsjailBackend("/usr/local/bin/nsjail");
-    await backend.exec("print(1)");
-
-    expect(stdinWritten).toBe("");
-  });
-
-  it("returns data injection error when stdin write fails with data", async () => {
-    let killed = false;
-    spawnResult = {
-      stdin: {
-        write: () => { throw new Error("EPIPE: broken pipe"); },
-        end: () => {},
-      },
-      stdout: makeStream(""),
-      stderr: makeStream(""),
-      exited: Promise.resolve(1),
-      kill: () => { killed = true; },
-    };
-
-    const backend = createPythonNsjailBackend("/usr/local/bin/nsjail");
-    const result = await backend.exec("print(df)", { columns: ["a"], rows: [[1]] });
-
-    expect(result.success).toBe(false);
-    if (!result.success) {
-      expect(result.error).toContain("Failed to inject data");
-    }
-    expect(killed).toBe(true);
-  });
-
-  it("continues execution when stdin write fails without data", async () => {
-    spawnResult = {
-      stdin: {
-        write: () => { throw new Error("EPIPE"); },
-        end: () => {},
-      },
-      stdout: makeStream(""),
-      stderr: makeStream("Python execution failed"),
-      exited: Promise.resolve(1),
-      kill: () => {},
-    };
-
-    const backend = createPythonNsjailBackend("/usr/local/bin/nsjail");
-    const result = await backend.exec("print(1)");
-
-    // Should fall through to normal error handling, not the data injection error
-    expect(result.success).toBe(false);
-    if (!result.success) {
-      expect(result.error).not.toContain("inject data");
-    }
-  });
-
-  it("detects stdout truncation and reports it", async () => {
-    // Simulate output that is exactly MAX_OUTPUT (1MB) — truncated
-    const bigOutput = "x".repeat(1024 * 1024);
-    setSpawnResult(bigOutput, "", 0);
-
-    const backend = createPythonNsjailBackend("/usr/local/bin/nsjail");
-    const result = await backend.exec("print('a' * 10000000)");
-
-    expect(result.success).toBe(false);
-    if (!result.success) {
-      expect(result.error).toContain("exceeded 1 MB");
-    }
-  });
-
-  it("returns unparseable output error for malformed result JSON", async () => {
-    const backend = createPythonNsjailBackend("/usr/local/bin/nsjail");
-
-    const savedSpawn = Bun.spawn;
-    Bun.spawn = ((...args: unknown[]) => {
-      spawnCalls.push({ args: [args[0]], options: args[1] });
-      const opts = args[1] as { env: Record<string, string> };
-      const marker = opts.env.ATLAS_RESULT_MARKER;
-      return {
-        stdin: { write: () => {}, end: () => {} },
-        stdout: makeStream(`${marker}NOT-VALID-JSON\n`),
-        stderr: makeStream("some error"),
-        exited: Promise.resolve(0),
-        kill: () => {},
-      };
-    }) as unknown as typeof Bun.spawn;
-
-    const result = await backend.exec("print(1)");
-
-    expect(result.success).toBe(false);
-    if (!result.success) {
-      expect(result.error).toContain("unparseable output");
-    }
-
-    Bun.spawn = savedSpawn;
-  });
-});
-
-describe("backend selection in python.ts", () => {
-  const savedEnv: Record<string, string | undefined> = {};
-
-  function saveAndSetEnv(key: string, value: string | undefined) {
-    if (!(key in savedEnv)) savedEnv[key] = process.env[key];
-    if (value === undefined) delete process.env[key];
-    else process.env[key] = value;
-  }
-
-  afterEach(() => {
-    for (const [key, value] of Object.entries(savedEnv)) {
-      if (value === undefined) delete process.env[key];
-      else process.env[key] = value;
-    }
-    for (const key of Object.keys(savedEnv)) delete savedEnv[key];
-  });
-
-  it("routes to Vercel sandbox backend when on Vercel", async () => {
-    saveAndSetEnv("ATLAS_SANDBOX_URL", undefined);
-    saveAndSetEnv("ATLAS_RUNTIME", "vercel");
-    saveAndSetEnv("ATLAS_NSJAIL_PATH", undefined);
-
-    // The AST import guard calls python3 — make it return valid output
-    const savedSpawn = Bun.spawn;
-    Bun.spawn = ((...args: unknown[]) => {
-      const cmd = args[0] as string[];
-      if (cmd[0] === "python3") {
-        return {
-          stdin: { write: () => {}, end: () => {} },
-          stdout: makeStream('{"imports":[],"calls":[]}'),
-          stderr: makeStream(""),
-          exited: Promise.resolve(0),
-          kill: () => {},
-        };
-      }
-      spawnCalls.push({ args: [args[0]], options: args[1] });
-      return spawnResult;
-    }) as unknown as typeof Bun.spawn;
-
-    const { executePython } = await import("@atlas/api/lib/tools/python");
-    const result = await executePython.execute!(
-      { code: 'print("hello")', explanation: "test", data: undefined },
-      {} as never,
-    ) as { success: boolean; error?: string };
-
-    // On Vercel, the backend tries to import @vercel/sandbox.
-    // In the test env it's not installed, so we get a sandbox-related error.
-    expect(result.success).toBe(false);
-    expect(result.error).toContain("sandbox");
-
-    Bun.spawn = savedSpawn;
-  });
-
-  it("returns hard-fail error when ATLAS_SANDBOX=nsjail but nsjail unavailable", async () => {
-    saveAndSetEnv("ATLAS_SANDBOX_URL", undefined);
-    saveAndSetEnv("ATLAS_RUNTIME", undefined);
-    saveAndSetEnv("VERCEL", undefined);
-    saveAndSetEnv("ATLAS_SANDBOX", "nsjail");
-
-    // Mock explore-nsjail to report nsjail as unavailable
-    const savedSpawn = Bun.spawn;
-    Bun.spawn = ((...args: unknown[]) => {
-      const cmd = args[0] as string[];
-      // The import guard calls python3 — make it return valid output
-      if (cmd[0] === "python3") {
-        return {
-          stdin: { write: () => {}, end: () => {} },
-          stdout: makeStream('{"imports":[],"calls":[]}'),
-          stderr: makeStream(""),
-          exited: Promise.resolve(0),
-          kill: () => {},
-        };
-      }
-      spawnCalls.push({ args: [args[0]], options: args[1] });
-      return spawnResult;
-    }) as unknown as typeof Bun.spawn;
-
-    // findNsjailBinary uses fs.accessSync which is mocked to always succeed,
-    // so nsjail will appear "found" and createPythonNsjailBackend will be called.
-    // The hard-fail test needs nsjail to NOT be found. Since fs is globally mocked
-    // to always succeed, we test the error message format via the Vercel test above
-    // and verify the hard-fail path exists in code review instead.
-    //
-    // Instead, test that when ATLAS_SANDBOX=nsjail AND a backend IS available,
-    // it uses it (doesn't skip to "no backend" error).
-    const { executePython } = await import("@atlas/api/lib/tools/python");
-    await executePython.execute!(
-      { code: 'print("hello")', explanation: "test", data: undefined },
-      {} as never,
-    );
-
-    // With mocked fs (accessSync always succeeds), nsjail appears available.
-    // The exec proceeds to the nsjail backend (which uses our mocked Bun.spawn).
-    // This verifies the ATLAS_SANDBOX=nsjail path routes to nsjail, not to error.
-    expect(spawnCalls.length).toBeGreaterThanOrEqual(1);
-    const nsjailCall = spawnCalls.find(c => {
-      const args = c.args[0] as string[];
-      return args[0]?.includes("nsjail");
-    });
-    expect(nsjailCall).toBeDefined();
-
-    Bun.spawn = savedSpawn;
-  });
-});