chore: handle test failures and appropriately mark them as failed in accuracyRunStatus

Author: himanshusinghs

scripts/mark-accuracy-run-finished.ts

@@ -1,7 +1,21 @@
 import { getAccuracySnapshotStorage } from "../tests/accuracy/sdk/accuracy-snapshot-storage/get-snapshot-storage.js";
+import {
+    AccuracyRunStatus,
+    AccuracyRunStatuses,
+} from "../tests/accuracy/sdk/accuracy-snapshot-storage/snapshot-storage.js";
 
-console.time(`Marked accuracy run id - ${process.env.MDB_ACCURACY_RUN_ID} as finished in`);
+let status: AccuracyRunStatuses | undefined;
+if (process.env.MDB_ACCURACY_RUN_STATUS === "done") {
+    status = AccuracyRunStatus.Done;
+} else if (process.env.MDB_ACCURACY_RUN_STATUS === "failed") {
+    status = AccuracyRunStatus.Failed;
+} else {
+    console.info(`Unknown status - ${process.env.MDB_ACCURACY_RUN_STATUS}, will not update accuracy run.`);
+    process.exit(1);
+}
+
+console.time(`Marked accuracy run id - ${process.env.MDB_ACCURACY_RUN_ID} as ${status} in`);
 const storage = await getAccuracySnapshotStorage();
-await storage.accuracyRunFinished();
+await storage.updateAccuracyRunStatus(status);
 await storage.close();
-console.timeEnd(`Marked accuracy run id - ${process.env.MDB_ACCURACY_RUN_ID} as finished in`);
+console.timeEnd(`Marked accuracy run id - ${process.env.MDB_ACCURACY_RUN_ID} as ${status} in`);

scripts/run-accuracy-tests.sh

@@ -22,7 +22,10 @@ node --experimental-vm-modules node_modules/jest/bin/jest.js --testPathPattern "
 
 # Each test run submits an accuracy snapshot entry for each prompt with the
 # accuracyRunStatus: "in-progress". When all the tests are done and jest exits
-# with an exit code of 0, we can safely mark accuracy run as finished.
+# with an exit code of 0, we can safely mark accuracy run as finished otherwise
+# failed.
 if [ $? -eq 0 ]; then
-  npx tsx scripts/mark-accuracy-run-finished.ts
+  MDB_ACCURACY_RUN_STATUS="done" npx tsx scripts/mark-accuracy-run-finished.ts 
+else
+  MDB_ACCURACY_RUN_STATUS="failed" npx tsx scripts/mark-accuracy-run-finished.ts
 fi

tests/accuracy/collection-schema.test.ts

@@ -1,34 +1,12 @@
 import { describeAccuracyTests, describeSuite } from "./sdk/describe-accuracy-tests.js";
 import { getAvailableModels } from "./sdk/models.js";
 import { AccuracyTestConfig } from "./sdk/describe-accuracy-tests.js";
-import { collectionSchemaResponse } from "../../src/tools/mongodb/metadata/collectionSchema.js";
-import { getSimplifiedSchema } from "mongodb-schema";
 
 function callsCollectionSchema(prompt: string): AccuracyTestConfig {
     return {
         injectConnectedAssumption: true,
         prompt: prompt,
-        mockedTools: {
-            "collection-schema": async function collectionSchema() {
-                return collectionSchemaResponse(
-                    "db1",
-                    "coll1",
-                    await getSimplifiedSchema([
-                        {
-                            name: "Sample name1",
-                            dob: "28.11.2001",
-                            location: "NY",
-                        },
-                        {
-                            name: "Sample name1",
-                            dob: "28.11.2001",
-                            location: "NY",
-                            title: "Dr.",
-                        },
-                    ])
-                );
-            },
-        },
+        mockedTools: {},
         expectedToolCalls: [
             {
                 toolName: "collection-schema",

tests/accuracy/sdk/accuracy-snapshot-storage/disk-snapshot-storage.ts

@@ -3,6 +3,7 @@ import fs from "fs/promises";
 import { fileURLToPath } from "url";
 import {
     AccuracyRunStatus,
+    AccuracyRunStatuses,
     AccuracySnapshotEntry,
     AccuracySnapshotEntrySchema,
     AccuracySnapshotStorage,
@@ -57,13 +58,13 @@ export class DiskSnapshotStorage implements AccuracySnapshotStorage {
         return latestRunId ? snapshot.filter((entry) => entry.accuracyRunId === latestRunId) : [];
     }
 
-    async accuracyRunFinished(): Promise<void> {
+    async updateAccuracyRunStatus(status: AccuracyRunStatuses) {
         const snapshot = await this.readSnapshot();
         const updatedSnapshot = snapshot.map((entry) => {
             if (entry.accuracyRunId === this.accuracyRunId) {
                 return {
                     ...entry,
-                    accuracyRunStatus: AccuracyRunStatus.Done,
+                    accuracyRunStatus: status,
                 };
             }
 

tests/accuracy/sdk/accuracy-snapshot-storage/mdb-snapshot-storage.ts

@@ -1,6 +1,7 @@
 import { Collection, MongoClient } from "mongodb";
 import {
     AccuracyRunStatus,
+    AccuracyRunStatuses,
     AccuracySnapshotEntry,
     AccuracySnapshotEntrySchema,
     AccuracySnapshotStorage,
@@ -76,10 +77,10 @@ export class MongoDBSnapshotStorage implements AccuracySnapshotStorage {
         return AccuracySnapshotEntrySchema.array().parse(snapshotEntries);
     }
 
-    async accuracyRunFinished(): Promise<void> {
+    async updateAccuracyRunStatus(status: AccuracyRunStatuses) {
         await this.snapshotCollection.updateMany(
             { accuracyRunId: this.accuracyRunId },
-            { $set: { accuracyRunStatus: AccuracyRunStatus.Done } }
+            { $set: { accuracyRunStatus: status } }
         );
     }
 

tests/accuracy/sdk/accuracy-snapshot-storage/snapshot-storage.ts

@@ -11,14 +11,17 @@ export type ActualToolCall = z.infer<typeof ActualToolCallSchema>;
 
 export const AccuracyRunStatus = {
     Done: "done",
+    Failed: "failed",
     InProgress: "in-progress",
 } as const;
 
+export type AccuracyRunStatuses = (typeof AccuracyRunStatus)[keyof typeof AccuracyRunStatus];
+
 export const AccuracySnapshotEntrySchema = z.object({
     // Git and meta information for snapshot entries
     accuracyRunId: z.string(),
     accuracyRunStatus: z
-        .enum([AccuracyRunStatus.Done, AccuracyRunStatus.InProgress])
+        .enum([AccuracyRunStatus.Done, AccuracyRunStatus.Failed, AccuracyRunStatus.InProgress])
         .default(AccuracyRunStatus.InProgress),
     createdOn: z.number(),
     commitSHA: z.string(),
@@ -67,7 +70,7 @@ export interface AccuracySnapshotStorage {
 
     getLatestSnapshotsForCommit(commit: string): Promise<AccuracySnapshotEntry[]>;
 
-    accuracyRunFinished(): Promise<void>;
+    updateAccuracyRunStatus(status: AccuracyRunStatuses): Promise<void>;
 
     close(): Promise<void>;
 }