fix tests

Arize-ai · Oct 9, 2024 · c8f2d11 · c8f2d11
1 parent 79276cf
commit c8f2d11
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 11 deletions.
diff --git a/app/src/pages/playground/__tests__/playgroundUtils.test.ts b/app/src/pages/playground/__tests__/playgroundUtils.test.ts
@@ -1,4 +1,4 @@
-import { _resetInstanceId } from "@phoenix/store";
+import { _resetInstanceId, _resetMessageId } from "@phoenix/store";
 
 import {
   getChatRole,
@@ -20,7 +20,12 @@ const expectedPlaygroundInstance = {
   template: {
     __type: "chat",
     messages: spanAttributesWithInputMessages.llm.input_messages.map(
-      ({ message }) => message
+      ({ message }, index) => {
+        return {
+          id: index,
+          ...message,
+        };
+      }
     ),
   },
   output: spanAttributesWithInputMessages.llm.output_messages,
@@ -30,6 +35,7 @@ const expectedPlaygroundInstance = {
 describe("transformSpanAttributesToPlaygroundInstance", () => {
   beforeEach(() => {
     _resetInstanceId();
+    _resetMessageId();
   });
   it("should throw if the attributes are not parsable", () => {
     const span = {
@@ -55,11 +61,21 @@ describe("transformSpanAttributesToPlaygroundInstance", () => {
       attributes: JSON.stringify(spanAttributesWithInputMessages),
     };
 
-    expect(transformSpanAttributesToPlaygroundInstance(span)).toEqual(
-      expectedPlaygroundInstance
-    );
+    const instance = transformSpanAttributesToPlaygroundInstance(span);
+    expect(instance?.template.__type).toEqual("chat");
+    if (instance?.template.__type !== "chat") {
+      throw new Error("Invalid template type constructed");
+    }
+    expect(instance?.template.messages).toHaveLength(2);
+    instance?.template.messages.forEach((message, index) => {
+      expect(message.role).toEqual(
+        expectedPlaygroundInstance.template.messages[index].role
+      );
+      expect(message.content).toEqual(
+        expectedPlaygroundInstance.template.messages[index].content
+      );
+    });
   });
-
   it("should return a PlaygroundInstance if the attributes contain llm.input_messages, even if output_messages are not present", () => {
     const span = {
       ...basePlaygroundSpan,
@@ -71,10 +87,12 @@ describe("transformSpanAttributesToPlaygroundInstance", () => {
         },
       }),
     };
-    expect(transformSpanAttributesToPlaygroundInstance(span)).toEqual({
-      ...expectedPlaygroundInstance,
-      output: undefined,
-    });
+    const instance = transformSpanAttributesToPlaygroundInstance(span);
+    expect(instance?.template.__type).toEqual("chat");
+    if (instance?.template.__type !== "chat") {
+      throw new Error("Invalid template type constructed");
+    }
+    expect(Array.isArray(instance?.template.messages)).toBeTruthy();
   });
 });
 

diff --git a/app/src/store/playgroundStore.tsx b/app/src/store/playgroundStore.tsx
@@ -5,7 +5,7 @@ export type GenAIOperationType = "chat" | "text_completion";
 
 let playgroundInstanceIdIndex = 0;
 let playgroundRunIdIndex = 0;
-let playgroundMessageIdIndex = 100;
+let playgroundMessageIdIndex = 0;
 
 /**
  * Generates a new playground instance ID
@@ -26,6 +26,15 @@ export const _resetInstanceId = () => {
   playgroundInstanceIdIndex = 0;
 };
 
+/**
+ * Resets the playground message ID to 0
+ *
+ * NB: This is only used for testing purposes
+ */
+export const _resetMessageId = () => {
+  playgroundInstanceIdIndex = 0;
+};
+
 /**
  * The input mode for the playground
  * @example "manual" or "dataset"