Fix mock checking in integration tests, wrapSDK and sporadic flakiness (

#636)

js/src/tests/run_trees.int.test.ts

@@ -4,7 +4,12 @@ import {
   RunTreeConfig,
   convertToDottedOrderFormat,
 } from "../run_trees.js";
-import { toArray, waitUntil, pollRunsUntilCount } from "./utils.js";
+import {
+  toArray,
+  waitUntil,
+  pollRunsUntilCount,
+  sanitizePresignedUrls,
+} from "./utils.js";
 
 test.concurrent(
   "Test post and patch run",
@@ -111,11 +116,15 @@ test.concurrent(
     const traceRuns = await toArray(traceRunsIter);
     expect(traceRuns.length).toEqual(5);
     // Sort by dotted order and assert runs lists are equal
-    const sortedRuns = runs.sort((a, b) =>
-      (a?.dotted_order ?? "").localeCompare(b?.dotted_order ?? "")
+    const sortedRuns = sanitizePresignedUrls(
+      runs.sort((a, b) =>
+        (a?.dotted_order ?? "").localeCompare(b?.dotted_order ?? "")
+      )
     );
-    const sortedTraceRuns = traceRuns.sort((a, b) =>
-      (a?.dotted_order ?? "").localeCompare(b?.dotted_order ?? "")
+    const sortedTraceRuns = sanitizePresignedUrls(
+      traceRuns.sort((a, b) =>
+        (a?.dotted_order ?? "").localeCompare(b?.dotted_order ?? "")
+      )
     );
     expect(sortedRuns).toEqual(sortedTraceRuns);
     await langchainClient.deleteProject({ projectName });

js/src/tests/utils.ts

@@ -96,3 +96,19 @@ export async function waitUntilRunFound(
     5_000
   );
 }
+
+export function sanitizePresignedUrls(payload: unknown) {
+  return JSON.parse(JSON.stringify(payload), (key, value) => {
+    if (key === "presigned_url") {
+      try {
+        const url = new URL(value);
+        url.searchParams.set("Signature", "[SIGNATURE]");
+        url.searchParams.set("Expires", "[EXPIRES]");
+        return url.toString();
+      } catch {
+        return value;
+      }
+    }
+    return value;
+  });
+}

js/src/tests/wrapped_openai.int.test.ts

@@ -6,7 +6,7 @@ import { wrapOpenAI } from "../wrappers/index.js";
 import { Client } from "../client.js";
 
 test.concurrent("chat.completions", async () => {
-  const client = new Client();
+  const client = new Client({ autoBatchTracing: false });
   const callSpy = jest
     .spyOn((client as any).caller, "call")
     .mockResolvedValue({ ok: true, text: () => "" });
@@ -62,8 +62,9 @@ test.concurrent("chat.completions", async () => {
   }
 
   expect(patchedChoices).toEqual(originalChoices);
+  expect(callSpy.mock.calls.length).toBeGreaterThanOrEqual(1);
   for (const call of callSpy.mock.calls) {
-    expect((call[2] as any)["method"]).toBe("POST");
+    expect(["POST", "PATCH"]).toContain((call[2] as any)["method"]);
   }
 
   const patchedStream2 = await patchedClient.chat.completions.create(
@@ -92,13 +93,14 @@ test.concurrent("chat.completions", async () => {
   }
 
   expect(patchedChoices2).toEqual(originalChoices);
+  expect(callSpy.mock.calls.length).toBeGreaterThanOrEqual(1);
   for (const call of callSpy.mock.calls) {
-    expect((call[2] as any)["method"]).toBe("POST");
+    expect(["POST", "PATCH"]).toContain((call[2] as any)["method"]);
   }
 });
 
 test.concurrent("chat completions with tool calling", async () => {
-  const client = new Client();
+  const client = new Client({ autoBatchTracing: false });
   const callSpy = jest
     // eslint-disable-next-line @typescript-eslint/no-explicit-any
     .spyOn((client as any).caller, "call")
@@ -229,9 +231,10 @@ test.concurrent("chat completions with tool calling", async () => {
   expect(removeToolCallId(patchedChoices)).toEqual(
     removeToolCallId(originalChoices)
   );
+  expect(callSpy.mock.calls.length).toBeGreaterThanOrEqual(1);
   for (const call of callSpy.mock.calls) {
     // eslint-disable-next-line @typescript-eslint/no-explicit-any
-    expect((call[2] as any)["method"]).toBe("POST");
+    expect(["POST", "PATCH"]).toContain((call[2] as any)["method"]);
   }
 
   const patchedStream2 = await patchedClient.chat.completions.create(
@@ -268,14 +271,15 @@ test.concurrent("chat completions with tool calling", async () => {
   expect(removeToolCallId(patchedChoices2)).toEqual(
     removeToolCallId(originalChoices)
   );
+  expect(callSpy.mock.calls.length).toBeGreaterThanOrEqual(1);
   for (const call of callSpy.mock.calls) {
     // eslint-disable-next-line @typescript-eslint/no-explicit-any
-    expect((call[2] as any)["method"]).toBe("POST");
+    expect(["POST", "PATCH"]).toContain((call[2] as any)["method"]);
   }
 });
 
 test.concurrent("completions", async () => {
-  const client = new Client();
+  const client = new Client({ autoBatchTracing: false });
   const callSpy = jest
     .spyOn((client as any).caller, "call")
     .mockResolvedValue({ ok: true, text: () => "" });
@@ -336,8 +340,9 @@ test.concurrent("completions", async () => {
   }
 
   expect(patchedChoices).toEqual(originalChoices);
+  expect(callSpy.mock.calls.length).toBeGreaterThanOrEqual(1);
   for (const call of callSpy.mock.calls) {
-    expect((call[2] as any)["method"]).toBe("POST");
+    expect(["POST", "PATCH"]).toContain((call[2] as any)["method"]);
   }
 
   const patchedStream2 = await patchedClient.completions.create(
@@ -365,9 +370,10 @@ test.concurrent("completions", async () => {
   }
 
   expect(patchedChoices2).toEqual(originalChoices);
+  expect(callSpy.mock.calls.length).toBeGreaterThanOrEqual(1);
   for (const call of callSpy.mock.calls) {
     // eslint-disable-next-line @typescript-eslint/no-explicit-any
-    expect((call[2] as any)["method"]).toBe("POST");
+    expect(["POST", "PATCH"]).toContain((call[2] as any)["method"]);
   }
 });
 

js/src/tests/wrapped_sdk.int.test.ts

@@ -4,7 +4,7 @@ import { wrapSDK } from "../wrappers/index.js";
 import { Client } from "../client.js";
 
 test.concurrent("chat.completions", async () => {
-  const client = new Client();
+  const client = new Client({ autoBatchTracing: false });
   const callSpy = jest
     // eslint-disable-next-line @typescript-eslint/no-explicit-any
     .spyOn((client as any).caller, "call")
@@ -58,8 +58,9 @@ test.concurrent("chat.completions", async () => {
   }
 
   expect(patchedChoices).toEqual(originalChoices);
+  expect(callSpy.mock.calls.length).toBeGreaterThanOrEqual(1);
   for (const call of callSpy.mock.calls) {
     // eslint-disable-next-line @typescript-eslint/no-explicit-any
-    expect((call[2] as any)["method"]).toBe("POST");
+    expect(["POST", "PATCH"]).toContain((call[2] as any)["method"]);
   }
 });

js/src/wrappers/openai.ts

@@ -282,7 +282,7 @@ const _wrapClient = <T extends object>(
           originalValue.bind(target),
           Object.assign(
             { name: [runName, propKey.toString()].join("."), run_type: "llm" },
-            options?.client
+            options
           )
         );
       } else if (