From d9b71b0f97d90e36a8906c463e0fa494ead96b87 Mon Sep 17 00:00:00 2001
From: Charlie Lamb <charlielamb20@icloud.com>
Date: Tue, 21 Apr 2026 11:41:29 +0100
Subject: [PATCH] chore: fix integration tests

---
 server/src/db/initDrizzle.ts            | 30 ++++++++++++++++++++-----
 server/src/utils/envUtils.ts            |  4 ++--
 server/tests/setup-integration-tests.ts |  2 +-
 3 files changed, 27 insertions(+), 9 deletions(-)

diff --git a/server/src/db/initDrizzle.ts b/server/src/db/initDrizzle.ts
index abae8fdc6..97292e8af 100644
--- a/server/src/db/initDrizzle.ts
+++ b/server/src/db/initDrizzle.ts
@@ -24,6 +24,19 @@ const normalizeExecuteRows = <TRow>(result: unknown): TRow[] => {
 	return result as TRow[];
 };
 
+const normalizeDbExecute = <
+	TDb extends { execute: (query: string | SQLWrapper) => Promise<unknown> },
+>(
+	db: TDb,
+) => {
+	const execute = db.execute.bind(db);
+	return Object.assign(db, {
+		execute: async <TRow = Record<string, unknown>>(
+			query: string | SQLWrapper,
+		) => normalizeExecuteRows<TRow>(await execute(query)),
+	});
+};
+
 /** Creates a Drizzle pool with the given configuration. */
 export const initDrizzle = ({
 	maxConnections = 10,
@@ -50,12 +63,17 @@ export const initDrizzle = ({
 	});
 
 	const drizzleDb = drizzle(client, { schema });
-	const execute = drizzleDb.execute.bind(drizzleDb);
-	const db = Object.assign(drizzleDb, {
-		execute: async <TRow = Record<string, unknown>>(
-			query: string | SQLWrapper,
-		) => normalizeExecuteRows<TRow>(await execute(query)),
-	}) as unknown as AutumnDb;
+	const transaction = drizzleDb.transaction.bind(drizzleDb);
+	const db = normalizeDbExecute(drizzleDb) as unknown as AutumnDb;
+	const normalizedTransaction: typeof drizzleDb.transaction = ((
+		fn,
+		config,
+	) =>
+		transaction(
+			(tx) => fn(normalizeDbExecute(tx) as typeof tx),
+			config,
+		)) as typeof drizzleDb.transaction;
+	db.transaction = normalizedTransaction as typeof db.transaction;
 
 	if (otelConfig.drizzle) {
 		instrumentDrizzleClient(db);
diff --git a/server/src/utils/envUtils.ts b/server/src/utils/envUtils.ts
index 422bf1193..d4a31cd09 100644
--- a/server/src/utils/envUtils.ts
+++ b/server/src/utils/envUtils.ts
@@ -4,8 +4,8 @@ import { config } from "dotenv";
 let hasLoadedLocalEnv = false;
 const shouldLogLocalEnvLoading = false;
 
-export const loadLocalEnv = () => {
-	if (hasLoadedLocalEnv) return;
+export const loadLocalEnv = ({ force = false }: { force?: boolean } = {}) => {
+	if (hasLoadedLocalEnv && !force) return;
 	hasLoadedLocalEnv = true;
 
 	const processDir = process.cwd();
diff --git a/server/tests/setup-integration-tests.ts b/server/tests/setup-integration-tests.ts
index 6736366e1..ab3c5a85a 100644
--- a/server/tests/setup-integration-tests.ts
+++ b/server/tests/setup-integration-tests.ts
@@ -32,6 +32,6 @@ if (isUnitTest()) {
 } else {
 	console.log("--- Setup integration tests ---");
 	await loadInfisicalSecrets();
-	loadLocalEnv();
+	loadLocalEnv({ force: true });
 	console.log("--- Setup integration tests complete ---");
 }