diff --git a/packages/app-tests/e2e/document-flow/stepper-component.spec.ts b/packages/app-tests/e2e/document-flow/stepper-component.spec.ts
index 6bab9248c..ec58b5f29 100644
--- a/packages/app-tests/e2e/document-flow/stepper-component.spec.ts
+++ b/packages/app-tests/e2e/document-flow/stepper-component.spec.ts
@@ -632,10 +632,10 @@ test('[DOCUMENT_FLOW]: should be able to create and sign a document with 3 recip
   }
 
   // Wait for the document to be signed.
-  await page.waitForTimeout(5000);
-
-  const signedDocument = await getDocumentById({ id: document.id, userId: user.id });
-  expect(signedDocument?.status).toBe(DocumentStatus.COMPLETED);
+  await expect(async () => {
+    const signedDocument = await getDocumentById({ id: document.id, userId: user.id });
+    expect(signedDocument?.status).toBe(DocumentStatus.COMPLETED);
+  }).toPass();
 });
 
 test('[DOCUMENT_FLOW]: should prevent out-of-order signing in sequential mode', async ({
@@ -747,8 +747,8 @@ test('[DOCUMENT_FLOW]: should be able to self sign a document', async ({ page })
   const updatedRecipient = await getRecipientById({ documentId: document.id, id: recipientId });
   expect(updatedRecipient?.signingStatus).toBe(SigningStatus.SIGNED);
 
-  await page.waitForTimeout(5000);
-
-  const signedDocument = await getDocumentById({ id: document.id, userId: user.id });
-  expect(signedDocument?.status).toBe(DocumentStatus.COMPLETED);
+  await expect(async () => {
+    const signedDocument = await getDocumentById({ id: document.id, userId: user.id });
+    expect(signedDocument?.status).toBe(DocumentStatus.COMPLETED);
+  }).toPass();
 });
diff --git a/packages/lib/server-only/document/self-sign-document.ts b/packages/lib/server-only/document/self-sign-document.ts
index 8e1c95220..1926bd9f0 100644
--- a/packages/lib/server-only/document/self-sign-document.ts
+++ b/packages/lib/server-only/document/self-sign-document.ts
@@ -31,7 +31,7 @@ export const selfSignDocument = async ({
     },
   });
 
-  const document = await prisma.document.findUnique({
+  const document = await prisma.document.findFirstOrThrow({
     where: {
       id: documentId,
       userId,
@@ -64,7 +64,7 @@ export const selfSignDocument = async ({
 
   const { documentData } = document;
 
-  if (!documentData.data) {
+  if (!documentData || !documentData.data) {
     throw new Error('Document data not found');
   }
 
diff --git a/packages/lib/server-only/document/send-document.tsx b/packages/lib/server-only/document/send-document.tsx
index 93ec12108..ada44718a 100644
--- a/packages/lib/server-only/document/send-document.tsx
+++ b/packages/lib/server-only/document/send-document.tsx
@@ -100,7 +100,7 @@ export const sendDocument = async ({
 
   const { documentData } = document;
 
-  if (!documentData.data) {
+  if (!documentData || !documentData.data) {
     throw new Error('Document data not found');
   }