2

I am trying to mock the @databricks/sql package using the jest.mock function and it continually pulls from the node modules package instead of the code I pass into the function.

My test file:

import mockResponse from "./mockResponse.json";

const {
  handler,
} = require("../../../backend/function/dataLakeConnect/src/index");

const testEvent = {
  arguments: {
    page: 1,
    clientId: "test",
    startTime: "2022-09-26 00:00:00",
    endTime: "2022-09-27 00:00:00",
  },
};

const exampleResponse = {
  ARN: "x",
  Name: "credentials/databricks/read-write",
  VersionId: "x",
  SecretString:
    '{"DATABRICKS_TOKEN":"test","DATABRICKS_SERVER_HOSTNAME":"test.test.com","DATABRICKS_HTTP_PATH":"/pathy/pathpath/mcgee"}',
  VersionStages: ["x"],
  CreatedDate: "x",
};

jest.mock("@databricks/sql", () => {
  return {
    DBSQLClient: function () {
      return {
        connect: function () {
          return {
            promise: function () {},
          };
        },
        openSession: function () {
          console.log("openSession called");
          return {
            promise: function () {
              return {
                executeStatement: function (sql, params) {
                  return {
                    promise: function () {
                      return {
                        fetchAll: function () {
                          return {
                            promise: function () {
                              return mockResponse;
                            },
                          };
                        },
                      };
                    },
                  };
                },
              };
            },
          };
        },
      };
    },
  };
});

jest.mock("aws-sdk", () => {
  return {
    SecretsManager: function () {
      return {
        getSecretValue: function ({ SecretId }) {
          if (SecretId === "credentials/databricks/read-write") {
            return {
              promise: function () {
                return exampleResponse;
              },
            };
          } else {
            throw new Error("mock error");
          }
        },
      };
    },
  };
});


describe("dataLakeConnect", () => {
  test("should return expected array", async () => {
    let res = await handler(testEvent);
    expect(res).toEqual(mockResponse);
  });
});

The function I am trying to test:

const AWS = require("aws-sdk");
const { DBSQLClient } = require("@databricks/sql");
const client = new AWS.SecretsManager({
  region: "us-east-1",
});
console.log(AWS);
console.log(DBSQLClient);
const databricksClient = new DBSQLClient();
const resultsPerPage = 300;

const getSecret = async (secretName) => {
  try {
    let data = await client.getSecretValue({ SecretId: secretName }).promise();
    if (data.SecretString) {
      return JSON.parse(data.SecretString);
    }
  } catch (err) {
    throw err;
  }
};

exports.handler = async (event) => {
  console.log(`EVENT: ${JSON.stringify(event)}`);
  const { page, clientId, startTime, endTime } = event.arguments;
  const offset = resultsPerPage * (page - 1);
  let sql = `SELECT * FROM test`;

  try {
    const creds = await getSecret("credentials/databricks/read-write");
    // console.log(databricksClient);
    await databricksClient.connect({
      token: creds.DATABRICKS_TOKEN,
      host: creds.DATABRICKS_SERVER_HOSTNAME,
      path: creds.DATABRICKS_HTTP_PATH,
    });
    const session = await databricksClient.openSession();
    const queryOperation = await session.executeStatement(sql, {
      runAsync: true,
      maxRows: 10000,
    });
    const result = await queryOperation.fetchAll({
      progress: false,
      callback: () => {},
    });

    await queryOperation.close();
    await session.close();
    await databricksClient.close();
    return result;
  } catch (err) {
    throw err;
  }
};

I am expecting the console logs in this file to return the mocked code, and the aws-sdk mock seems to be working fine, however, the @databricks/sql console log returns the original module and this causes the test to timeout.

I have tried putting these mocks a __mocks__ directory next to the node modules like jest suggests, but I end up with errors saying jest cannot find the package.

Any suggestions on how to actually get the @databricks/sql mock to work? I am pretty stumped.

  • How complex those test doubles have to be should suggest that this is something you _don't_ want to be mocking. Don't mock what you don't own - test this kind of thing at the integration level. – jonrsharpe Sep 27 '22 at 17:01
  • 1
    Try moving your `jest.mock` statements to the first thing in the test file. I think because you are requiring the handler before the jest mocks, the handler doesn't know about the mocks – lucasvw Sep 27 '22 at 17:03
  • on top of the jest.mock that should be above all imports importing what you mock, do your async functions mocks work? I would have expected just some nested async functions and a final Promise resolve: openSession: async () => ({executeStatement: async () => ({fetchAll: () => Promise.resolve(mockResponse)})}) – user3252327 Sep 27 '22 at 17:15

1 Answers1

1

With your code the mock actually (copied the code from question) work but I get:

dataLakeConnect › should return expected array

    TypeError: session.executeStatement is not a function

      35 |     });
      36 |     const session = await databricksClient.openSession();
    > 37 |     const queryOperation = await session.executeStatement(sql, {
         |                                          ^
      38 |       runAsync: true,
      39 |       maxRows: 10000,
      40 |     });

After changing the mock to:

jest.mock("@databricks/sql", () => ({
  DBSQLClient: function () {
    return {
      connect: () => Promise.resolve({}),
      openSession: () => {
        console.log("openSession called");
        return Promise.resolve({
          executeStatement: () =>
            Promise.resolve({
              fetchAll: () => Promise.resolve(mockResponse),
              close: () => Promise.resolve({}),
            }),
          close: () => Promise.resolve({}),
        });
      },
      close: () => Promise.resolve({}),
    };
  },
}));

The test passes:

 ● Console

    console.log
      { SecretsManager: [Function: SecretsManager] }

      at Object.<anonymous> (src/totest.js:6:9)

    console.log
      [Function: DBSQLClient]

      at Object.<anonymous> (src/totest.js:7:9)

    console.log
      EVENT: {"arguments":{"page":1,"clientId":"test","startTime":"2022-09-26 00:00:00","endTime":"2022-09-27 00:00:00"}}

      at handler (src/totest.js:23:11)

    console.log
      openSession called

      at Object.openSession (src/__tests__/totest.spec.js:54:17)
Marek Rozmus
  • 773
  • 7
  • 13