diff --git a/tests/catalyst.test.mts b/tests/catalyst.test.mts
index 201a4e0..686665c 100644
--- a/tests/catalyst.test.mts
+++ b/tests/catalyst.test.mts
@@ -3,19 +3,21 @@ import { describe, it, expect, vi, beforeEach } from 'vitest';
 // Mock dependencies
 
 vi.mock('../src/puml/EntityParser.mjs', () => ({
-  EntityParser: vi.fn().mockImplementation(() => ({
-    parse: vi.fn().mockReturnValue([]),
-    getObjectWithPropertyAndValueInHierarchy: vi.fn().mockReturnValue(null)
-  }))
+  EntityParser: class MockEntityParser {
+    parse = vi.fn().mockReturnValue([])
+    getObjectWithPropertyAndValueInHierarchy = vi.fn().mockReturnValue(null)
+  }
 }));
 
 vi.mock('../src/mx/Mx.mjs', () => ({
-  Mx: vi.fn().mockImplementation(() => ({
-    addMxC4: vi.fn(),
-    addMxC4Relationship: vi.fn(),
-    generate: vi.fn().mockResolvedValue('<xml>test</xml>')
-  })),
-  MxGeometry: vi.fn()
+  Mx: class MockMx {
+    addMxC4 = vi.fn()
+    addMxC4Relationship = vi.fn()
+    generate = vi.fn().mockResolvedValue('<xml>test</xml>')
+  },
+  MxGeometry: class MockMxGeometry {
+    constructor(x: number, y: number, width: number, height: number) {}
+  }
 }));
 
 vi.mock('../src/puml/RelParser.mjs', () => ({
diff --git a/tests/mx/Mx.test.mts b/tests/mx/Mx.test.mts
index eabca57..5c332bd 100644
--- a/tests/mx/Mx.test.mts
+++ b/tests/mx/Mx.test.mts
@@ -4,9 +4,9 @@ import { Mx, MxGeometry } from '../../src/mx/Mx.mjs';
 // Mock xml2js
 vi.mock('xml2js', () => ({
   default: {
-    Builder: vi.fn().mockImplementation(() => ({
-      buildObject: vi.fn().mockReturnValue('<xml>test</xml>')
-    }))
+    Builder: class MockBuilder {
+      buildObject = vi.fn().mockReturnValue('<xml>test</xml>')
+    }
   }
 }));