diff --git a/packages/cli/src/modules/data-store/__tests__/data-store.service.test.ts b/packages/cli/src/modules/data-store/__tests__/data-store.service.test.ts
index 7dbc574eb2e..df98df51992 100644
--- a/packages/cli/src/modules/data-store/__tests__/data-store.service.test.ts
+++ b/packages/cli/src/modules/data-store/__tests__/data-store.service.test.ts
@@ -13,9 +13,7 @@ beforeAll(async () => {
 });
 
 beforeEach(async () => {
-	try {
-		await testDb.truncate(['DataStoreEntity']);
-	} catch {}
+	await testDb.truncate(['DataStoreEntity', 'DataStoreColumnEntity']);
 });
 
 afterAll(async () => {
@@ -52,7 +50,7 @@ describe('dataStore', () => {
 	});
 
 	afterEach(async () => {
-		// this kinda sucks
+		// Clean up any created user data stores
 		await dataStoreService.deleteDataStoreAll();
 	});
 
diff --git a/packages/cli/test/integration/shared/types.ts b/packages/cli/test/integration/shared/types.ts
index 290afe5a53b..ad20066445b 100644
--- a/packages/cli/test/integration/shared/types.ts
+++ b/packages/cli/test/integration/shared/types.ts
@@ -42,7 +42,8 @@ type EndpointGroup =
 	| 'evaluation'
 	| 'ai'
 	| 'folder'
-	| 'insights';
+	| 'insights'
+	| 'data-store';
 
 type ModuleName = 'insights' | 'external-secrets' | 'data-store';