87 lines · 3.2 KB · TypeScript
RawDownload
import { expect, test } from "bun:test";import { chmodSync, mkdtempSync, readFileSync, rmSync, writeFileSync } from "node:fs";import { tmpdir } from "node:os";import { join } from "node:path";import type { GpuInfo } from "../../models/types";import { queryNvidiaSmiSnapshot } from "./gpu"; test("collects stable NVIDIA identities from one shared snapshot", async () => {  const directory = mkdtempSync(join(tmpdir(), "local-studio-nvidia-smi-"));  const binaryPath = join(directory, "nvidia-smi");  const argumentsPath = join(directory, "arguments");  const countPath = join(directory, "count");  const previousBinary = process.env["NVIDIA_SMI_PATH"];  const previousArgumentsPath = process.env["LOCAL_STUDIO_NVIDIA_SMI_ARGUMENTS_PATH"];  const previousCountPath = process.env["LOCAL_STUDIO_NVIDIA_SMI_COUNT_PATH"];   try {    writeFileSync(      binaryPath,      [        "#!/bin/sh",        'printf \'%s\\n\' "$@" > "$LOCAL_STUDIO_NVIDIA_SMI_ARGUMENTS_PATH"',        "printf '1\\n' >> \"$LOCAL_STUDIO_NVIDIA_SMI_COUNT_PATH\"",        "printf '%s\\n' 'GPU-3090, 00000000:82:00.0, NVIDIA GeForce RTX 3090, 24576, 2048, 22528, 12, 48, 80.5, 350, 580.95.05'",      ].join("\n"),    );    chmodSync(binaryPath, 0o755);    process.env["NVIDIA_SMI_PATH"] = binaryPath;    process.env["LOCAL_STUDIO_NVIDIA_SMI_ARGUMENTS_PATH"] = argumentsPath;    process.env["LOCAL_STUDIO_NVIDIA_SMI_COUNT_PATH"] = countPath;     const snapshot = await queryNvidiaSmiSnapshot();     expect(snapshot).toEqual({      available: true,      driverVersion: "580.95.05",      gpus: [        {          uuid: "GPU-3090",          pci_bus_id: "00000000:82:00.0",          index: 0,          name: "NVIDIA GeForce RTX 3090",          memory_total_mb: 24576,          memory_used_mb: 2048,          memory_free_mb: 22528,          utilization_pct: 12,          temp_c: 48,          power_draw: 80.5,          power_limit: 350,        },      ],    });    expect(readFileSync(countPath, "utf8")).toBe("1\n");    expect(readFileSync(argumentsPath, "utf8").trim().split("\n")).toEqual([      "--query-gpu=uuid,pci.bus_id,name,memory.total,memory.used,memory.free,utilization.gpu,temperature.gpu,power.draw,power.limit,driver_version",      "--format=csv,noheader,nounits",    ]);  } finally {    if (previousBinary === undefined) delete process.env["NVIDIA_SMI_PATH"];    else process.env["NVIDIA_SMI_PATH"] = previousBinary;    if (previousArgumentsPath === undefined) {      delete process.env["LOCAL_STUDIO_NVIDIA_SMI_ARGUMENTS_PATH"];    } else {      process.env["LOCAL_STUDIO_NVIDIA_SMI_ARGUMENTS_PATH"] = previousArgumentsPath;    }    if (previousCountPath === undefined) delete process.env["LOCAL_STUDIO_NVIDIA_SMI_COUNT_PATH"];    else process.env["LOCAL_STUDIO_NVIDIA_SMI_COUNT_PATH"] = previousCountPath;    rmSync(directory, { recursive: true, force: true });  }}); test("keeps stable identity optional for non-NVIDIA telemetry", () => {  const gpu: GpuInfo = {    index: 0,    name: "Intel Arc Pro B70",    memory_total_mb: 24576,    memory_used_mb: 0,    memory_free_mb: 24576,    utilization_pct: 0,    temp_c: 0,    power_draw: 0,    power_limit: 0,  };   expect(gpu.uuid).toBeUndefined();  expect(gpu.pci_bus_id).toBeUndefined();});