From 1dc96e16258194fe943d6ad11ea791e7b063f4fc Mon Sep 17 00:00:00 2001
From: Damien BUTY <dam.buty@gmail.com>
Date: Mon, 1 Jul 2024 17:13:47 +0200
Subject: [PATCH] fix(types): getXXX methods return instances

---
 src/api.ts                    | 33 ++++++++++++++++++++-------------
 tests/integration/api.test.ts |  4 ++--
 2 files changed, 22 insertions(+), 15 deletions(-)

diff --git a/src/api.ts b/src/api.ts
index 4fb0ece..4d2befb 100644
--- a/src/api.ts
+++ b/src/api.ts
@@ -32,8 +32,8 @@ import {
   Prompt,
   Score,
   Step,
-  StepConstructor,
   StepType,
+  Thread,
   User,
   Utils
 } from './types';
@@ -467,7 +467,7 @@ export class API {
     before?: Maybe<string>;
     filters?: StepsFilter[];
     orderBy?: StepsOrderBy;
-  }): Promise<PaginatedResponse<OmitUtils<StepConstructor>>> {
+  }): Promise<PaginatedResponse<Step>> {
     const query = `
       query GetSteps(
         $after: ID,
@@ -509,7 +509,7 @@ export class API {
 
     const response = result.data.steps;
 
-    response.data = response.edges.map((x: any) => x.node);
+    response.data = response.edges.map((x: any) => new Step(this, x.node));
     delete response.edges;
 
     return response;
@@ -537,9 +537,11 @@ export class API {
 
     const result = await this.makeGqlCall(query, variables);
 
-    const step = result.data.step;
+    if (!result.data.step) {
+      return null;
+    }
 
-    return step;
+    return new Step(this, result.data.step);
   }
 
   /**
@@ -876,7 +878,7 @@ export class API {
     };
 
     const response = await this.makeGqlCall(query, variables);
-    return response.data.upsertThread;
+    return new Thread(this, response.data.upsertThread);
   }
 
   /**
@@ -897,7 +899,7 @@ export class API {
     filters?: ThreadsFilter[];
     orderBy?: ThreadsOrderBy;
     stepTypesToKeep?: StepType[];
-  }): Promise<PaginatedResponse<CleanThreadFields>> {
+  }): Promise<PaginatedResponse<Thread>> {
     const query = `
     query GetThreads(
         $after: ID,
@@ -941,7 +943,7 @@ export class API {
 
     const response = result.data.threads;
 
-    response.data = response.edges.map((x: any) => x.node);
+    response.data = response.edges.map((x: any) => new Thread(this, x.node));
     delete response.edges;
 
     return response;
@@ -953,7 +955,7 @@ export class API {
    * @param id - The unique identifier of the thread. This parameter is required.
    * @returns The detailed information of the specified thread.
    */
-  async getThread(id: string) {
+  async getThread(id: string): Promise<Maybe<Thread>> {
     const query = `
     query GetThread($id: String!) {
         threadDetail(id: $id) {
@@ -965,7 +967,12 @@ export class API {
     const variables = { id };
 
     const response = await this.makeGqlCall(query, variables);
-    return response.data.threadDetail;
+
+    if (!response.data.threadDetail) {
+      return null;
+    }
+
+    return new Thread(this, response.data.threadDetail);
   }
 
   /**
@@ -974,7 +981,7 @@ export class API {
    * @param id - The unique identifier of the thread to be deleted. This parameter is required.
    * @returns The ID of the deleted thread.
    */
-  async deleteThread(id: string) {
+  async deleteThread(id: string): Promise<string> {
     const query = `
     mutation DeleteThread($threadId: String!) {
         deleteThread(id: $threadId) {
@@ -1051,7 +1058,7 @@ export class API {
 
     const response = result.data.participants;
 
-    response.data = response.edges.map((x: any) => x.node);
+    response.data = response.edges.map((x: any) => new User(x.node));
     delete response.edges;
 
     return response;
@@ -1273,7 +1280,7 @@ export class API {
 
     const response = result.data.scores;
 
-    response.data = response.edges.map((x: any) => x.node);
+    response.data = response.edges.map((x: any) => new Score(x.node));
     delete response.edges;
 
     return response;
diff --git a/tests/integration/api.test.ts b/tests/integration/api.test.ts
index ad68e63..fd4a218 100644
--- a/tests/integration/api.test.ts
+++ b/tests/integration/api.test.ts
@@ -82,7 +82,7 @@ describe('End to end tests for the SDK', function () {
     expect(thread.metadata).toStrictEqual({ foo: 'bar' });
 
     const fetchedThread = await client.api.getThread(thread.id);
-    expect(fetchedThread.id).toBe(thread.id);
+    expect(fetchedThread?.id).toBe(thread.id);
 
     const updatedThread = await client.api.upsertThread({
       threadId: thread.id,
@@ -112,7 +112,7 @@ describe('End to end tests for the SDK', function () {
     expect(thread.metadata).toStrictEqual({ foo: 'bar' });
 
     const fetchedThread = await client.api.getThread(thread.id);
-    expect(fetchedThread.id).toBe(thread.id);
+    expect(fetchedThread?.id).toBe(thread.id);
 
     const updatedThread = await client.api.upsertThread(
       thread.id,