Mocking chained functions with Jest

I like to make fun of things, but mocking functions is not my favourite pastime.

Too soon ? Oh well...

Don't worry, I promise, no more puns in this article. Aaah-wait I lied.

Seriously tho, were here to talk about testing and the difficulties when testing more than an example function from a tutorial.

I've always found that writing tests is easy for standalone functions.

However, it becomes disproportionately hard as soon as you try to test a function that uses a third-party library and requires mocking/stubbing, especially with chained functions/methods.

Now, let's discuss how to mock the Google Cloud SDK's BigQuery Node.JS API using their sample code:

bigquery.dataset(datasetId).table(tableId).insert(data);

Note: Although it may seem, upon reading this, that I figured things out quickly, on the contrary, the process involved a few hours of head-banging, reading articles, documentation, source codes, StackOverflow, and nagging ChatGPT and my colleague JP.

The task

My task at hand was quite simple: insert a single row into a BigQuery table. I grabbed the example and quickly set my function up then went on to write some tests, covering the scenarios that can arise when the my function is called.

The stuff that didn't work

Initially, I thought just mocking the class with Jest will work just fine, but things kept erroring.

// Mock the BigQuery class
jest.mock('@google-cloud/bigquery');

Later, I switched to mock implementation, but it still didn't fully work. It passed for toHaveBeenCalledWith for the first chained function of .table() but not for the subsequent ones.

This got me thinking and made me look at the type signatures and implementation of the bigquery package itself in my node modules.

The realisation

Then, the realisation struck.

The return type of bigquery.dataset class included a line indicating that it returns a function called table.

table(id: string, options?: TableOptions): Table;

Upon further investigation, I found a similar pattern. Each function returns a class containing the next function I was trying to use, so my mock implementation needed to account for this.

The solution

The function

// Import the Google Cloud client library
const {BigQuery} = require('@google-cloud/bigquery'); 
const bigQueryClient = new BigQuery();
 
// Inserts the JSON objects into my_dataset:my_table.
async function insertRowsAsStream(data) {
const datasetId = 'my_dataset';
const tableId = 'my_table';
const rows = [
  {name: 'Tom', age: 30},
  {name: 'Jane', age: 32},
];
 
// Insert data into a table
await bigQueryClient
  .dataset(datasetId)
  .table(tableId)
  .insert(data);
console.log(`Inserted ${data.length} row(s)`);}
};
 
module.exports = { bigQueryClient, insertRowsAsStream}

The working test

// Import our method and our configured BigQuery client
const { bigQueryClient, insertRowsAsStream } = require('.myModule');
 
// Mock the BigQuery class and its methods
jest.mock('@google-cloud/bigquery', () => {
const mockTableInstance = {
  insert: jest.fn(),
};
 
const mockDatasetInstance = {
  table: jest.fn(() => mockTableInstance),
};
 
const mockBigQueryInstance = {
  dataset: jest.fn(() => mockDatasetInstance),
};
return { BigQuery: jest.fn(() => mockBigQueryInstance) };
});
 
// Test cases for insertRowsAsStream function
describe('insertRowsAsStream', () => {
it('should insert data into BigQuery and log success message', async () => {
  // Mock data
  const data = { name: 'Ben', age: 33 };
  const datasetId = 'testDatasetId';
  const tableId = 'testTableId';
 
  await insertRowsAsStream(data, tableId);
 
  // Assert statements
  expect(bigquery.dataset).toHaveBeenCalledWith(dataSet);
  expect(bigquery.dataset().table).toHaveBeenCalledWith(tableId);
  expect(bigquery.dataset().table().insert).toHaveBeenCalledWith({
    name: data.name,
    age: data.age,
  });
});
});