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.