diff --git a/tests/machines.getById.test.ts b/tests/machines.getById.test.ts
index 66e0140..6d44f69 100644
--- a/tests/machines.getById.test.ts
+++ b/tests/machines.getById.test.ts
@@ -55,18 +55,32 @@ describe("convex.machines.getById", () => {
         if (id === machine._id) return machine
         return null
       }),
-      query: vi.fn((_table: string) => ({
-        withIndex: vi.fn((_name: string, _cb: unknown) => ({
-          collect: vi.fn(async () => [
-            { revoked: false, expiresAt: FIXED_NOW + 60_000, lastUsedAt: FIXED_NOW - 1000, usageCount: 5 },
-          ]),
-          take: vi.fn(async () => [
-            { revoked: false, expiresAt: FIXED_NOW + 60_000, lastUsedAt: FIXED_NOW - 1000, usageCount: 5 },
-          ]),
-        })),
-        collect: vi.fn(async () => []),
-        take: vi.fn(async () => []),
-      })),
+      query: vi.fn((table: string) => {
+        // Mock para machineHeartbeats - retorna o lastHeartbeatAt da nova tabela
+        if (table === "machineHeartbeats") {
+          return {
+            withIndex: vi.fn((_name: string, _cb: unknown) => ({
+              first: vi.fn(async () => ({ machineId: machine._id, lastHeartbeatAt: FIXED_NOW - 1000 })),
+              collect: vi.fn(async () => [{ machineId: machine._id, lastHeartbeatAt: FIXED_NOW - 1000 }]),
+              take: vi.fn(async () => [{ machineId: machine._id, lastHeartbeatAt: FIXED_NOW - 1000 }]),
+            })),
+          }
+        }
+        // Mock para machineTokens e outras tabelas
+        return {
+          withIndex: vi.fn((_name: string, _cb: unknown) => ({
+            first: vi.fn(async () => null),
+            collect: vi.fn(async () => [
+              { revoked: false, expiresAt: FIXED_NOW + 60_000, lastUsedAt: FIXED_NOW - 1000, usageCount: 5 },
+            ]),
+            take: vi.fn(async () => [
+              { revoked: false, expiresAt: FIXED_NOW + 60_000, lastUsedAt: FIXED_NOW - 1000, usageCount: 5 },
+            ]),
+          })),
+          collect: vi.fn(async () => []),
+          take: vi.fn(async () => []),
+        }
+      }),
     }
 
     const ctx = { db } as unknown as Parameters<typeof getByIdHandler>[0]