diff --git a/client/tests/e2e/suite/extension.e2e.ts b/client/tests/e2e/suite/extension.e2e.ts
index 9be0d51..527cd1e 100644
--- a/client/tests/e2e/suite/extension.e2e.ts
+++ b/client/tests/e2e/suite/extension.e2e.ts
@@ -23,7 +23,7 @@ suite("Extension Test Suite", () => {
             const docUri = getDocUri(path.join("test_tool_01.xml"));
             await activateAndOpenInEditor(docUri);
 
-            await waitForDiagnostics();
+            await waitForDiagnostics(docUri);
             await assertDiagnostics(docUri, []);
         });
 
@@ -31,7 +31,7 @@ suite("Extension Test Suite", () => {
             const docUri = getDocUri(path.join("macros_01.xml"));
             await activateAndOpenInEditor(docUri);
 
-            await waitForDiagnostics();
+            await waitForDiagnostics(docUri);
             await assertDiagnostics(docUri, []);
         });
 
@@ -39,7 +39,7 @@ suite("Extension Test Suite", () => {
             const docUri = getDocUri(path.join("test_invalid_tool_01.xml"));
             await activateAndOpenInEditor(docUri);
 
-            await waitForDiagnostics();
+            await waitForDiagnostics(docUri);
             await assertDiagnostics(docUri, [
                 {
                     message: "Element 'tool': The attribute 'id' is required but missing.",
@@ -58,7 +58,7 @@ suite("Extension Test Suite", () => {
             const docUri = getDocUri(path.join("test_syntax_error_01.xml"));
             await activateAndOpenInEditor(docUri);
 
-            await waitForDiagnostics();
+            await waitForDiagnostics(docUri);
             await assertDiagnostics(docUri, [
                 {
                     message: "Couldn't find end of Start Tag tool line 1, line 2, column 1",
@@ -72,7 +72,7 @@ suite("Extension Test Suite", () => {
             const docUri = getDocUri(path.join("test_syntax_error_macro_01.xml"));
             await activateAndOpenInEditor(docUri);
 
-            await waitForDiagnostics();
+            await waitForDiagnostics(docUri);
             await assertDiagnostics(docUri, [
                 {
                     message: "Premature end of data in tag macros line 1, line 2, column 1",
diff --git a/client/tests/e2e/suite/helpers.ts b/client/tests/e2e/suite/helpers.ts
index ea3f162..59c2b9c 100644
--- a/client/tests/e2e/suite/helpers.ts
+++ b/client/tests/e2e/suite/helpers.ts
@@ -47,8 +47,15 @@ export async function sleep(ms: number): Promise<void> {
     return new Promise((resolve) => setTimeout(resolve, ms));
 }
 
-export async function waitForDiagnostics(): Promise<void> {
-    return sleep(500);
+export async function waitForDiagnostics(docUri: vscode.Uri, timeoutInMilliseconds: number = 2000): Promise<void> {
+    const waitMilliseconds = 100;
+    let waitTimeout = timeoutInMilliseconds;
+    let diagnostics = vscode.languages.getDiagnostics(docUri);
+    while (waitTimeout > 0 && !diagnostics.length) {
+        await sleep(waitMilliseconds);
+        waitTimeout -= waitMilliseconds;
+        diagnostics = vscode.languages.getDiagnostics(docUri);
+    }
 }
 
 export const getDocPath = (filePath: string): string => {