From 5eb5fcd17ca401168a6f0a0b805b48c29ae13a3e Mon Sep 17 00:00:00 2001
From: Felix Mosheev <9304194+felixmosh@users.noreply.github.com>
Date: Tue, 17 Aug 2021 10:03:13 +0300
Subject: [PATCH] Improve support for `first` & `pluck` methods, closes #3 (#4)

---
 src/MockClient.ts      |  1 +
 src/Tracker.ts         | 11 ++++++++++-
 tests/select.spec.ts   | 24 +++++++++++++++++++-----
 types/mock-client.d.ts |  2 ++
 4 files changed, 32 insertions(+), 6 deletions(-)

diff --git a/src/MockClient.ts b/src/MockClient.ts
index 01789c2..beaef04 100644
--- a/src/MockClient.ts
+++ b/src/MockClient.ts
@@ -31,6 +31,7 @@ export class MockClient extends knex.Client {
     switch (rawMethod) {
       case 'first':
       case 'pluck':
+        rawQuery.postOp = rawMethod;
         method = 'select';
         break;
       case 'del':
diff --git a/src/Tracker.ts b/src/Tracker.ts
index a5ed60e..37eb9f1 100644
--- a/src/Tracker.ts
+++ b/src/Tracker.ts
@@ -63,7 +63,7 @@ export class Tracker {
               } else {
                 const data =
                   typeof handler.data === 'function' ? await handler.data(rawQuery) : handler.data;
-                resolve(cloneDeep(data));
+                resolve(cloneDeep(Tracker.applyPostOp(data, rawQuery)));
               }
 
               if (handler.once) {
@@ -135,4 +135,13 @@ export class Tracker {
       };
     };
   }
+
+  private static applyPostOp(data: any, rawQuery: RawQuery) {
+    if (rawQuery.postOp === 'first' && Array.isArray(data)) {
+      return data[0];
+    } else if (rawQuery.postOp === 'pluck' && rawQuery.pluck && Array.isArray(data)) {
+      return data.map((item) => item[rawQuery.pluck as string]);
+    }
+    return data;
+  }
 }
diff --git a/tests/select.spec.ts b/tests/select.spec.ts
index d949d3f..6900861 100644
--- a/tests/select.spec.ts
+++ b/tests/select.spec.ts
@@ -163,19 +163,33 @@ describe('mock Select statement', () => {
   });
 
   it('should support `first` query', async () => {
-    tracker.on.select('table_name').response([]);
+    const data = [
+      { id: faker.datatype.number({ min: 1 }) },
+      { id: faker.datatype.number({ min: 1 }) },
+    ];
 
-    await db('table_name').first();
+    tracker.on.select('table_name').response(data);
 
-    expect(tracker.history.select).toHaveLength(1);
+    const firstValue = await db('table_name').first();
+    const allValues = await db('table_name');
+
+    expect(tracker.history.select).toHaveLength(2);
+    expect(firstValue).toEqual(data[0]);
+    expect(allValues).toEqual(data);
   });
 
   it('should support `pluck` query', async () => {
-    tracker.on.select('table_name').response([]);
+    const data = [
+      { id: faker.datatype.number({ min: 1 }) },
+      { id: faker.datatype.number({ min: 1 }) },
+    ];
+
+    tracker.on.select('table_name').response(data);
 
-    await db('table_name').pluck('id');
+    const values = await db('table_name').pluck('id');
 
     expect(tracker.history.select).toHaveLength(1);
+    expect(values).toEqual(data.map(({ id }) => id));
   });
 
   it('should support `raw` select statement', async () => {
diff --git a/types/mock-client.d.ts b/types/mock-client.d.ts
index 2237044..4e13b8f 100644
--- a/types/mock-client.d.ts
+++ b/types/mock-client.d.ts
@@ -1,6 +1,8 @@
 export interface RawQuery {
   method: 'select' | 'insert' | 'update' | 'delete' | 'any';
   sql: string;
+  postOp?: 'first' | 'pluck';
+  pluck?: string;
   bindings: any[];
   options: Record<string, any>;
   timeout: boolean;