openlattice-ssh 0.0.3

package/tests/ssh-provider.test.ts

@@ -0,0 +1,778 @@
+import { describe, it, expect, vi, beforeEach } from "vitest";
+import { SSHProvider } from "../src/ssh-provider";
+
+// ── Mock ssh2 ───────────────────────────────────────────────────────
+
+const mocks = vi.hoisted(() => {
+  const mockStream = {
+    on: vi.fn(),
+    stderr: { on: vi.fn() },
+  };
+
+  const mockSftp = {
+    readFile: vi.fn(),
+    writeFile: vi.fn(),
+    readdir: vi.fn(),
+    unlink: vi.fn(),
+    mkdir: vi.fn(),
+  };
+
+  const mockExec = vi.fn();
+  const mockSftpFn = vi.fn();
+  const mockConnect = vi.fn();
+  const mockEnd = vi.fn();
+  const mockOn = vi.fn();
+
+  return {
+    mockStream,
+    mockSftp,
+    mockExec,
+    mockSftpFn,
+    mockConnect,
+    mockEnd,
+    mockOn,
+  };
+});
+
+vi.mock("ssh2", () => ({
+  Client: vi.fn(function (this: any) {
+    this.exec = mocks.mockExec;
+    this.sftp = mocks.mockSftpFn;
+    this.connect = mocks.mockConnect;
+    this.end = mocks.mockEnd;
+    this.on = mocks.mockOn;
+    return this;
+  }),
+}));
+
+// Helper to set up mockExec to call the callback with a mock stream
+function setupExec(stdout: string, exitCode = 0) {
+  mocks.mockExec.mockImplementation((_cmd: string, cb: Function) => {
+    const stream = {
+      on: vi.fn().mockImplementation(function (
+        this: any,
+        event: string,
+        handler: Function
+      ) {
+        if (event === "data") {
+          setTimeout(() => handler(Buffer.from(stdout)), 0);
+        }
+        if (event === "close") {
+          setTimeout(() => handler(exitCode), 5);
+        }
+        return this;
+      }),
+      stderr: {
+        on: vi.fn().mockReturnThis(),
+      },
+    };
+    cb(undefined, stream);
+  });
+}
+
+function setupExecWithStderr(
+  stdout: string,
+  stderr: string,
+  exitCode = 0
+) {
+  mocks.mockExec.mockImplementation((_cmd: string, cb: Function) => {
+    const stream = {
+      on: vi.fn().mockImplementation(function (
+        this: any,
+        event: string,
+        handler: Function
+      ) {
+        if (event === "data") {
+          setTimeout(() => handler(Buffer.from(stdout)), 0);
+        }
+        if (event === "close") {
+          setTimeout(() => handler(exitCode), 5);
+        }
+        return this;
+      }),
+      stderr: {
+        on: vi.fn().mockImplementation(function (
+          this: any,
+          event: string,
+          handler: Function
+        ) {
+          if (event === "data") {
+            setTimeout(() => handler(Buffer.from(stderr)), 0);
+          }
+          return this;
+        }),
+      },
+    };
+    cb(undefined, stream);
+  });
+}
+
+// ── Tests ───────────────────────────────────────────────────────────
+
+describe("SSHProvider", () => {
+  let provider: SSHProvider;
+
+  const defaultConfig = {
+    hosts: [
+      { host: "192.168.1.100", port: 22, username: "testuser" },
+    ],
+    defaultUser: "defaultuser",
+  };
+
+  beforeEach(() => {
+    vi.clearAllMocks();
+    // Reset the mock to fire "ready" on connect
+    mocks.mockOn.mockImplementation(function (
+      this: any,
+      event: string,
+      cb: Function
+    ) {
+      if (event === "ready") {
+        setTimeout(() => cb(), 0);
+      }
+      return this;
+    });
+    setupExec("", 0);
+    provider = new SSHProvider(defaultConfig);
+  });
+
+  describe("constructor", () => {
+    it("sets name to 'ssh'", () => {
+      expect(provider.name).toBe("ssh");
+    });
+
+    it("declares expected capabilities", () => {
+      expect(provider.capabilities.restart).toBe(true);
+      expect(provider.capabilities.pause).toBe(false);
+      expect(provider.capabilities.snapshot).toBe(false);
+      expect(provider.capabilities.logs).toBe(false);
+      expect(provider.capabilities.tailscale).toBe(true);
+      expect(provider.capabilities.persistentStorage).toBe(true);
+      expect(provider.capabilities.architectures).toContain("x86_64");
+      expect(provider.capabilities.architectures).toContain("arm64");
+    });
+
+    it("detects GPU from host config", () => {
+      const gpuProvider = new SSHProvider({
+        hosts: [{ host: "gpu-host", gpuAvailable: true }],
+      });
+      expect(gpuProvider.capabilities.gpu).toBe(true);
+    });
+
+    it("defaults gpu to false when no hosts have GPU", () => {
+      expect(provider.capabilities.gpu).toBe(false);
+    });
+
+    it("throws if no hosts configured", () => {
+      expect(() => new SSHProvider({ hosts: [] })).toThrow(
+        "at least one host"
+      );
+    });
+  });
+
+  describe("provision", () => {
+    it("creates working directory and returns node", async () => {
+      const node = await provider.provision({
+        runtime: { image: "any" },
+      });
+
+      expect(node.externalId).toMatch(/^192\.168\.1\.100:22\/ssh-/);
+      expect(node.endpoints).toHaveLength(1);
+      expect(node.endpoints[0].host).toBe("192.168.1.100");
+      expect(node.endpoints[0].port).toBe(22);
+    });
+
+    it("calls mkdir on the remote host", async () => {
+      await provider.provision({
+        runtime: { image: "any" },
+      });
+
+      expect(mocks.mockExec).toHaveBeenCalled();
+      const cmd = mocks.mockExec.mock.calls[0][0];
+      expect(cmd).toContain("mkdir -p /tmp/openlattice/");
+    });
+
+    it("runs initial command in background", async () => {
+      // First call: mkdir, second call: nohup
+      let callIdx = 0;
+      mocks.mockExec.mockImplementation((cmd: string, cb: Function) => {
+        const stdout = callIdx === 1 ? "1234\n" : "";
+        callIdx++;
+        const stream = {
+          on: vi.fn().mockImplementation(function (
+            this: any,
+            event: string,
+            handler: Function
+          ) {
+            if (event === "data") {
+              setTimeout(() => handler(Buffer.from(stdout)), 0);
+            }
+            if (event === "close") {
+              setTimeout(() => handler(0), 5);
+            }
+            return this;
+          }),
+          stderr: { on: vi.fn().mockReturnThis() },
+        };
+        cb(undefined, stream);
+      });
+
+      await provider.provision({
+        runtime: {
+          image: "any",
+          command: ["python", "server.py"],
+        },
+      });
+
+      expect(mocks.mockExec).toHaveBeenCalledTimes(2);
+      const nohupCmd = mocks.mockExec.mock.calls[1][0];
+      expect(nohupCmd).toContain("nohup python server.py");
+    });
+
+    it("rejects GPU spec when no GPU hosts", async () => {
+      await expect(
+        provider.provision({
+          runtime: { image: "any" },
+          gpu: { count: 1 },
+        })
+      ).rejects.toThrow("no hosts with GPU");
+    });
+
+    it("selects hosts by labels", async () => {
+      const multiHostProvider = new SSHProvider({
+        hosts: [
+          { host: "host-a", labels: { region: "us-east" } },
+          { host: "host-b", labels: { region: "eu-west" } },
+        ],
+      });
+
+      const node = await multiHostProvider.provision({
+        runtime: { image: "any" },
+        labels: { region: "eu-west" },
+      });
+
+      expect(node.externalId).toContain("host-b");
+    });
+
+    it("throws when no hosts match labels", async () => {
+      await expect(
+        provider.provision({
+          runtime: { image: "any" },
+          labels: { region: "unknown" },
+        })
+      ).rejects.toThrow("no hosts matching labels");
+    });
+
+    it("calls ensureTailscale when tailscale is true", async () => {
+      // Track exec calls to verify tailscale commands
+      const execCmds: string[] = [];
+      mocks.mockExec.mockImplementation((cmd: string, cb: Function) => {
+        execCmds.push(cmd);
+        const stdout = cmd.includes("which tailscale")
+          ? "/usr/bin/tailscale"
+          : cmd.includes("pgrep tailscaled")
+          ? ""
+          : "";
+        const stream = {
+          on: vi.fn().mockImplementation(function (
+            this: any,
+            event: string,
+            handler: Function
+          ) {
+            if (event === "data" && stdout) {
+              setTimeout(() => handler(Buffer.from(stdout)), 0);
+            }
+            if (event === "close") {
+              setTimeout(() => handler(0), 5);
+            }
+            return this;
+          }),
+          stderr: { on: vi.fn().mockReturnThis() },
+        };
+        cb(undefined, stream);
+      });
+
+      await provider.provision({
+        runtime: { image: "any" },
+        network: { tailscale: true },
+      });
+
+      // Should have called: mkdir, which tailscale, pgrep tailscaled
+      expect(execCmds.some((c) => c.includes("which tailscale"))).toBe(true);
+      expect(execCmds.some((c) => c.includes("pgrep tailscaled"))).toBe(true);
+    });
+
+    it("installs tailscale when not found on host", async () => {
+      const execCmds: string[] = [];
+      let callIdx = 0;
+      mocks.mockExec.mockImplementation((cmd: string, cb: Function) => {
+        execCmds.push(cmd);
+        const isWhichCheck = cmd.includes("which tailscale");
+        const exitCode = isWhichCheck && callIdx === 1 ? 1 : 0;
+        callIdx++;
+        const stream = {
+          on: vi.fn().mockImplementation(function (
+            this: any,
+            event: string,
+            handler: Function
+          ) {
+            if (event === "close") {
+              setTimeout(() => handler(exitCode), 5);
+            }
+            return this;
+          }),
+          stderr: { on: vi.fn().mockReturnThis() },
+        };
+        cb(undefined, stream);
+      });
+
+      await provider.provision({
+        runtime: { image: "any" },
+        network: { tailscale: true },
+      });
+
+      expect(
+        execCmds.some((c) => c.includes("tailscale.com/install.sh"))
+      ).toBe(true);
+    });
+
+    it("does not call ensureTailscale when tailscale is not set", async () => {
+      const execCmds: string[] = [];
+      mocks.mockExec.mockImplementation((cmd: string, cb: Function) => {
+        execCmds.push(cmd);
+        const stream = {
+          on: vi.fn().mockImplementation(function (
+            this: any,
+            event: string,
+            handler: Function
+          ) {
+            if (event === "close") {
+              setTimeout(() => handler(0), 5);
+            }
+            return this;
+          }),
+          stderr: { on: vi.fn().mockReturnThis() },
+        };
+        cb(undefined, stream);
+      });
+
+      await provider.provision({
+        runtime: { image: "any" },
+      });
+
+      // Should only have mkdir, no tailscale checks
+      expect(execCmds.some((c) => c.includes("tailscale"))).toBe(false);
+    });
+
+    it("runs tailscale up --authkey when tailscaleAuthKey is set", async () => {
+      const execCmds: string[] = [];
+      mocks.mockExec.mockImplementation((cmd: string, cb: Function) => {
+        execCmds.push(cmd);
+        const stdout = cmd.includes("which tailscale")
+          ? "/usr/bin/tailscale"
+          : "";
+        const stream = {
+          on: vi.fn().mockImplementation(function (
+            this: any,
+            event: string,
+            handler: Function
+          ) {
+            if (event === "data" && stdout) {
+              setTimeout(() => handler(Buffer.from(stdout)), 0);
+            }
+            if (event === "close") {
+              setTimeout(() => handler(0), 5);
+            }
+            return this;
+          }),
+          stderr: { on: vi.fn().mockReturnThis() },
+        };
+        cb(undefined, stream);
+      });
+
+      await provider.provision({
+        runtime: { image: "any" },
+        network: { tailscaleAuthKey: "tskey-auth-abc123" },
+      });
+
+      // Should have called tailscale up with the auth key
+      expect(
+        execCmds.some((c) => c.includes("tailscale up") && c.includes("tskey-auth-abc123"))
+      ).toBe(true);
+    });
+
+    it("enables tailscale setup when only tailscaleAuthKey is set", async () => {
+      const execCmds: string[] = [];
+      mocks.mockExec.mockImplementation((cmd: string, cb: Function) => {
+        execCmds.push(cmd);
+        const stdout = cmd.includes("which tailscale")
+          ? "/usr/bin/tailscale"
+          : "";
+        const stream = {
+          on: vi.fn().mockImplementation(function (
+            this: any,
+            event: string,
+            handler: Function
+          ) {
+            if (event === "data" && stdout) {
+              setTimeout(() => handler(Buffer.from(stdout)), 0);
+            }
+            if (event === "close") {
+              setTimeout(() => handler(0), 5);
+            }
+            return this;
+          }),
+          stderr: { on: vi.fn().mockReturnThis() },
+        };
+        cb(undefined, stream);
+      });
+
+      await provider.provision({
+        runtime: { image: "any" },
+        network: { tailscaleAuthKey: "tskey-auth-xyz" },
+      });
+
+      // Should have called ensureTailscale (which tailscale check)
+      expect(execCmds.some((c) => c.includes("which tailscale"))).toBe(true);
+      // And tailscale up
+      expect(
+        execCmds.some((c) => c.includes("tailscale up"))
+      ).toBe(true);
+    });
+  });
+
+  describe("exec", () => {
+    it("runs command on remote host", async () => {
+      setupExec("hello\n", 0);
+      const node = await provider.provision({
+        runtime: { image: "any" },
+      });
+
+      setupExec("world\n", 0);
+      const result = await provider.exec(node.externalId, [
+        "echo",
+        "world",
+      ]);
+
+      expect(result.exitCode).toBe(0);
+      expect(result.stdout).toBe("world\n");
+    });
+
+    it("captures stderr", async () => {
+      setupExec("", 0);
+      const node = await provider.provision({
+        runtime: { image: "any" },
+      });
+
+      setupExecWithStderr("", "err\n", 1);
+      const result = await provider.exec(node.externalId, [
+        "sh",
+        "-c",
+        "echo err >&2; exit 1",
+      ]);
+
+      expect(result.exitCode).toBe(1);
+      expect(result.stderr).toBe("err\n");
+    });
+
+    it("prepends cwd to command", async () => {
+      setupExec("", 0);
+      const node = await provider.provision({
+        runtime: { image: "any" },
+      });
+
+      setupExec("", 0);
+      await provider.exec(node.externalId, ["ls"], { cwd: "/app" });
+
+      const cmd = mocks.mockExec.mock.calls[mocks.mockExec.mock.calls.length - 1][0];
+      expect(cmd).toContain("cd /app");
+    });
+
+    it("prepends env vars to command", async () => {
+      setupExec("", 0);
+      const node = await provider.provision({
+        runtime: { image: "any" },
+      });
+
+      setupExec("", 0);
+      await provider.exec(node.externalId, ["env"], {
+        env: { FOO: "bar" },
+      });
+
+      const cmd = mocks.mockExec.mock.calls[mocks.mockExec.mock.calls.length - 1][0];
+      expect(cmd).toContain("FOO='bar'");
+    });
+
+    it("throws for unknown node", async () => {
+      await expect(
+        provider.exec("nonexistent", ["echo"])
+      ).rejects.toThrow("node not found");
+    });
+  });
+
+  describe("destroy", () => {
+    it("removes working directory", async () => {
+      setupExec("", 0);
+      const node = await provider.provision({
+        runtime: { image: "any" },
+      });
+
+      setupExec("", 0);
+      await provider.destroy(node.externalId);
+
+      // Check that rm -rf was called
+      const rmCmd = mocks.mockExec.mock.calls.find((c: string[]) =>
+        c[0].includes("rm -rf")
+      );
+      expect(rmCmd).toBeDefined();
+    });
+
+    it("is idempotent for already destroyed nodes", async () => {
+      await expect(
+        provider.destroy("nonexistent")
+      ).resolves.toBeUndefined();
+    });
+
+    it("removes node from internal state", async () => {
+      setupExec("", 0);
+      const node = await provider.provision({
+        runtime: { image: "any" },
+      });
+
+      setupExec("", 0);
+      await provider.destroy(node.externalId);
+
+      const status = await provider.inspect(node.externalId);
+      expect(status.status).toBe("terminated");
+    });
+  });
+
+  describe("inspect", () => {
+    it("returns running for active node", async () => {
+      setupExec("", 0);
+      const node = await provider.provision({
+        runtime: { image: "any" },
+      });
+
+      setupExec("ok\n", 0);
+      const status = await provider.inspect(node.externalId);
+      expect(status.status).toBe("running");
+      expect(status.startedAt).toBeInstanceOf(Date);
+    });
+
+    it("returns terminated for unknown node", async () => {
+      const status = await provider.inspect("nonexistent");
+      expect(status.status).toBe("terminated");
+    });
+
+    it("returns running/stopped based on PID check", async () => {
+      // Provision with a command so we get a PID
+      let callIdx = 0;
+      mocks.mockExec.mockImplementation((cmd: string, cb: Function) => {
+        const stdout = callIdx === 1 ? "99\n" : "";
+        callIdx++;
+        const stream = {
+          on: vi.fn().mockImplementation(function (
+            this: any,
+            event: string,
+            handler: Function
+          ) {
+            if (event === "data") {
+              setTimeout(() => handler(Buffer.from(stdout)), 0);
+            }
+            if (event === "close") {
+              setTimeout(() => handler(0), 5);
+            }
+            return this;
+          }),
+          stderr: { on: vi.fn().mockReturnThis() },
+        };
+        cb(undefined, stream);
+      });
+
+      const node = await provider.provision({
+        runtime: { image: "any", command: ["sleep", "100"] },
+      });
+
+      // Inspect - process running
+      setupExec("running\n", 0);
+      const status = await provider.inspect(node.externalId);
+      expect(status.status).toBe("running");
+
+      // Inspect - process stopped
+      setupExec("stopped\n", 0);
+      const status2 = await provider.inspect(node.externalId);
+      expect(status2.status).toBe("stopped");
+    });
+  });
+
+  describe("stop / start", () => {
+    it("stop kills the managed process", async () => {
+      // Provision with PID
+      let callIdx = 0;
+      mocks.mockExec.mockImplementation((_cmd: string, cb: Function) => {
+        const stdout = callIdx === 1 ? "42\n" : "";
+        callIdx++;
+        const stream = {
+          on: vi.fn().mockImplementation(function (
+            this: any,
+            event: string,
+            handler: Function
+          ) {
+            if (event === "data") {
+              setTimeout(() => handler(Buffer.from(stdout)), 0);
+            }
+            if (event === "close") {
+              setTimeout(() => handler(0), 5);
+            }
+            return this;
+          }),
+          stderr: { on: vi.fn().mockReturnThis() },
+        };
+        cb(undefined, stream);
+      });
+
+      const node = await provider.provision({
+        runtime: { image: "any", command: ["sleep", "100"] },
+      });
+
+      setupExec("", 0);
+      await provider.stop(node.externalId);
+
+      const killCmd = mocks.mockExec.mock.calls.find((c: string[]) =>
+        c[0].includes("kill 42")
+      );
+      expect(killCmd).toBeDefined();
+    });
+
+    it("start verifies connectivity", async () => {
+      setupExec("", 0);
+      const node = await provider.provision({
+        runtime: { image: "any" },
+      });
+
+      setupExec("ok\n", 0);
+      await expect(
+        provider.start(node.externalId)
+      ).resolves.toBeUndefined();
+    });
+  });
+
+  describe("healthCheck", () => {
+    it("returns healthy when all hosts reachable", async () => {
+      setupExec("ok\n", 0);
+      const health = await provider.healthCheck();
+
+      expect(health.healthy).toBe(true);
+      expect(health.latencyMs).toBeTypeOf("number");
+    });
+
+    it("returns unhealthy when host connection fails", async () => {
+      mocks.mockOn.mockImplementation(function (
+        this: any,
+        event: string,
+        cb: Function
+      ) {
+        if (event === "error") {
+          setTimeout(() => cb(new Error("connection refused")), 0);
+        }
+        return this;
+      });
+
+      // Need a fresh provider to trigger new connection
+      const freshProvider = new SSHProvider({
+        hosts: [{ host: "unreachable" }],
+      });
+
+      const health = await freshProvider.healthCheck();
+      expect(health.healthy).toBe(false);
+      expect(health.message).toContain("unreachable");
+    });
+  });
+
+  describe("getExtension", () => {
+    it("returns file extension", async () => {
+      setupExec("", 0);
+      const node = await provider.provision({
+        runtime: { image: "any" },
+      });
+
+      const ext = provider.getExtension(node.externalId, "files");
+      expect(ext).toBeDefined();
+      expect(ext).toHaveProperty("read");
+      expect(ext).toHaveProperty("write");
+      expect(ext).toHaveProperty("list");
+      expect(ext).toHaveProperty("remove");
+      expect(ext).toHaveProperty("mkdir");
+    });
+
+    it("returns network extension", async () => {
+      setupExec("", 0);
+      const node = await provider.provision({
+        runtime: { image: "any" },
+      });
+
+      const ext = provider.getExtension(node.externalId, "network");
+      expect(ext).toBeDefined();
+      expect(ext!.getUrl).toBeTypeOf("function");
+    });
+
+    it("network extension returns host URL", async () => {
+      setupExec("", 0);
+      const node = await provider.provision({
+        runtime: { image: "any" },
+      });
+
+      const ext = provider.getExtension(node.externalId, "network");
+      const url = await ext!.getUrl(8080);
+      expect(url).toBe("http://192.168.1.100:8080");
+    });
+
+    it("returns undefined for unsupported extensions", async () => {
+      setupExec("", 0);
+      const node = await provider.provision({
+        runtime: { image: "any" },
+      });
+
+      expect(
+        provider.getExtension(node.externalId, "metrics")
+      ).toBeUndefined();
+    });
+  });
+
+  describe("close", () => {
+    it("ends all connections", async () => {
+      setupExec("", 0);
+      await provider.provision({
+        runtime: { image: "any" },
+      });
+
+      await provider.close();
+      expect(mocks.mockEnd).toHaveBeenCalled();
+    });
+  });
+
+  describe("host selection", () => {
+    it("round-robins across hosts", async () => {
+      const multiProvider = new SSHProvider({
+        hosts: [
+          { host: "host-a" },
+          { host: "host-b" },
+        ],
+      });
+
+      setupExec("", 0);
+      const n1 = await multiProvider.provision({
+        runtime: { image: "any" },
+      });
+
+      setupExec("", 0);
+      const n2 = await multiProvider.provision({
+        runtime: { image: "any" },
+      });
+
+      expect(n1.externalId).toContain("host-a");
+      expect(n2.externalId).toContain("host-b");
+    });
+  });
+});