Tests: update pi SDK mocks

src/agents/model-catalog.test.ts

@@ -38,8 +38,12 @@ describe("loadModelCatalog", () => {
         throw new Error("boom");
       }
       return {
-        discoverAuthStorage: () => ({}),
-        discoverModels: () => [{ id: "gpt-4.1", name: "GPT-4.1", provider: "openai" }],
+        AuthStorage: class {},
+        ModelRegistry: class {
+          getAll() {
+            return [{ id: "gpt-4.1", name: "GPT-4.1", provider: "openai" }];
+          }
+        },
       } as unknown as PiSdkModule;
     });
 
@@ -59,19 +63,21 @@ describe("loadModelCatalog", () => {
     __setModelCatalogImportForTest(
       async () =>
         ({
-          discoverAuthStorage: () => ({}),
-          discoverModels: () => ({
-            getAll: () => [
-              { id: "gpt-4.1", name: "GPT-4.1", provider: "openai" },
-              {
-                get id() {
-                  throw new Error("boom");
+          AuthStorage: class {},
+          ModelRegistry: class {
+            getAll() {
+              return [
+                { id: "gpt-4.1", name: "GPT-4.1", provider: "openai" },
+                {
+                  get id() {
+                    throw new Error("boom");
+                  },
+                  provider: "openai",
+                  name: "bad",
                 },
-                provider: "openai",
-                name: "bad",
-              },
-            ],
-          }),
+              ];
+            }
+          },
         }) as unknown as PiSdkModule,
     );
 

src/agents/pi-embedded-runner/model.test.ts

@@ -1,8 +1,14 @@
 import { describe, expect, it, vi } from "vitest";
 
 vi.mock("@mariozechner/pi-coding-agent", () => ({
-  discoverAuthStorage: vi.fn(() => ({ mocked: true })),
-  discoverModels: vi.fn(() => ({ find: vi.fn(() => null) })),
+  AuthStorage: class {
+    mocked = true;
+  },
+  ModelRegistry: class {
+    find() {
+      return null;
+    }
+  },
 }));
 
 import type { OpenClawConfig } from "../../config/config.js";

src/commands/models.list.test.ts

@@ -13,8 +13,22 @@ const resolveProfileUnusableUntilForDisplay = vi.fn().mockReturnValue(null);
 const resolveEnvApiKey = vi.fn().mockReturnValue(undefined);
 const resolveAwsSdkEnvVarName = vi.fn().mockReturnValue(undefined);
 const getCustomProviderApiKey = vi.fn().mockReturnValue(undefined);
-const discoverAuthStorage = vi.fn().mockReturnValue({});
-const discoverModels = vi.fn();
+const modelRegistryState = {
+  models: [] as Array<Record<string, unknown>>,
+  available: [] as Array<Record<string, unknown>>,
+};
+
+class AuthStorage {}
+
+class ModelRegistry {
+  getAll() {
+    return modelRegistryState.models;
+  }
+
+  getAvailable() {
+    return modelRegistryState.available;
+  }
+}
 
 vi.mock("../config/config.js", () => ({
   CONFIG_PATH: "/tmp/openclaw.json",
@@ -45,8 +59,8 @@ vi.mock("../agents/model-auth.js", () => ({
 }));
 
 vi.mock("@mariozechner/pi-coding-agent", () => ({
-  discoverAuthStorage,
-  discoverModels,
+  AuthStorage,
+  ModelRegistry,
 }));
 
 function makeRuntime() {
@@ -99,10 +113,8 @@ describe("models list/status", () => {
       contextWindow: 128000,
     };
 
-    discoverModels.mockReturnValue({
-      getAll: () => [model],
-      getAvailable: () => [model],
-    });
+    modelRegistryState.models = [model];
+    modelRegistryState.available = [model];
 
     const { modelsListCommand } = await import("./models/list.js");
     await modelsListCommand({ json: true }, runtime);
@@ -127,10 +139,8 @@ describe("models list/status", () => {
       contextWindow: 128000,
     };
 
-    discoverModels.mockReturnValue({
-      getAll: () => [model],
-      getAvailable: () => [model],
-    });
+    modelRegistryState.models = [model];
+    modelRegistryState.available = [model];
 
     const { modelsListCommand } = await import("./models/list.js");
     await modelsListCommand({ plain: true }, runtime);
@@ -164,10 +174,8 @@ describe("models list/status", () => {
       },
     ];
 
-    discoverModels.mockReturnValue({
-      getAll: () => models,
-      getAvailable: () => models,
-    });
+    modelRegistryState.models = models;
+    modelRegistryState.available = models;
 
     const { modelsListCommand } = await import("./models/list.js");
     await modelsListCommand({ all: true, provider: "z.ai", json: true }, runtime);
@@ -203,10 +211,8 @@ describe("models list/status", () => {
       },
     ];
 
-    discoverModels.mockReturnValue({
-      getAll: () => models,
-      getAvailable: () => models,
-    });
+    modelRegistryState.models = models;
+    modelRegistryState.available = models;
 
     const { modelsListCommand } = await import("./models/list.js");
     await modelsListCommand({ all: true, provider: "Z.AI", json: true }, runtime);
@@ -242,10 +248,8 @@ describe("models list/status", () => {
       },
     ];
 
-    discoverModels.mockReturnValue({
-      getAll: () => models,
-      getAvailable: () => models,
-    });
+    modelRegistryState.models = models;
+    modelRegistryState.available = models;
 
     const { modelsListCommand } = await import("./models/list.js");
     await modelsListCommand({ all: true, provider: "z-ai", json: true }, runtime);
@@ -271,10 +275,8 @@ describe("models list/status", () => {
       contextWindow: 128000,
     };
 
-    discoverModels.mockReturnValue({
-      getAll: () => [model],
-      getAvailable: () => [],
-    });
+    modelRegistryState.models = [model];
+    modelRegistryState.available = [];
 
     const { modelsListCommand } = await import("./models/list.js");
     await modelsListCommand({ all: true, json: true }, runtime);