Improve support for first & pluck methods, closes #3 (#4)

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':

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;
+  }
 }

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 () => {

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;