Drizzle unit test (Vitest)

export const createNewLocationRow = async (
  createNewLocationPayload: ICreateLocation
) => {
  const { name, coordinate } = createNewLocationPayload;
  try {
    const result = await db
      .insert(locations)
      .values({
        name,
        coordinate,
      })
      .returning();
    return result[0];
  } catch (err) {
    console.log(err);
    throw new CustomError(
      "Failed to create new location",
      HTTP_CODES.BAD_GATEWAY
    );
  }
};

vi.mock("db", async (importOriginal) => {
  const pool = new Pool({
    user: "postgres",
    host: "localhost",
    database: "db2",
    password: "000000",
    port: 5432,
  });
  const db = drizzle(pool, { schema });
  return {
    ...(await importOriginal<typeof import("../../db")>()),
    db,
    pool,
  };
});

// // Apply migrations before each test
beforeEach(async () => {
  await applyMigrations();
});

// Clean up the database after each test
afterEach(async () => {
  await db.execute(sql`drop schema if exists public cascade`);
  await db.execute(sql`create schema public`);
  await db.execute(sql`drop schema if exists drizzle cascade`);
});

describe("Locations Repository", () => {
  describe("createNewLocationRow", () => {
    it("should create new location as expected", async () => {
      const mockQueryData: ICreateLocation = {
        name: "Testing",
        coordinate: [
          [-1, -1],
          [-5, -1],
          [-1, 2],
          [-5, 2],
        ],
      };

      const newLocation = await createNewLocationRow(mockQueryData);

      expect(newLocation.id).toEqual(1);
    });
  });
});