Rewrite logic using await

AikidoSec · Feb 16, 2024 · c0355c2 · c0355c2
1 parent 798aa18
commit c0355c2
Showing 1 changed file with 42 additions and 40 deletions.
diff --git a/end2end/tests/express-mongodb.test.js b/end2end/tests/express-mongodb.test.js
@@ -8,7 +8,11 @@ const pathToApp = resolve(
   "app.js"
 );
 
-t.test("it blocks in blocking mode", (t) => {
+async function timeout(ms) {
+  return new Promise((resolve) => setTimeout(resolve, ms));
+}
+
+t.test("it blocks in blocking mode", async () => {
   const server = spawn(`node`, [pathToApp], { shell: true });
 
   let stdout = "";
@@ -23,35 +27,34 @@ t.test("it blocks in blocking mode", (t) => {
     console.log("stderr", data.toString());
   });
 
-  server.unref();
+  // Wait for the server to start
+  await timeout(1000);
 
-  setTimeout(() => {
-    Promise.all([
+  try {
+    const [noSQLInjection, normalSearch] = await Promise.all([
       fetch("http://localhost:4000/?search[$ne]=null", {
         signal: AbortSignal.timeout(5000),
       }),
       fetch("http://localhost:4000/?search=title", {
         signal: AbortSignal.timeout(5000),
       }),
-    ])
-      .then((results) => {
-        const [noSQLInjection, normalSearch] = results;
-        t.equal(noSQLInjection.status, 500);
-        t.equal(normalSearch.status, 200);
-        t.match(stdout, /Starting agent/);
-        t.match(stderr, /Aikido guard has blocked a NoSQL injection/);
-      })
-      .catch((error) => {
-        t.fail(error.message);
-      })
-      .finally(() => {
-        server.kill();
-        t.end();
-      });
-  }, 1000);
+    ]);
+
+    t.equal(noSQLInjection.status, 500);
+    t.equal(normalSearch.status, 200);
+    t.match(stdout, /Starting agent/);
+    t.match(stderr, /Aikido guard has blocked a NoSQL injection/);
+  } catch (error) {
+    t.fail(error.message);
+  } finally {
+    await new Promise((resolve) => {
+      server.kill();
+      server.on("close", resolve);
+    });
+  }
 });
 
-t.test("it does not block in dry mode", (t) => {
+t.test("it does not block in dry mode", async () => {
   const server = spawn(`node`, [pathToApp], {
     env: { ...process.env, AIKIDO_NO_BLOCKING: "true" },
     shell: true,
@@ -69,30 +72,29 @@ t.test("it does not block in dry mode", (t) => {
     console.log("stderr", data.toString());
   });
 
-  server.unref();
+  // Wait for the server to start
+  await timeout(1000);
 
-  setTimeout(() => {
-    Promise.all([
+  try {
+    const [noSQLInjection, normalSearch] = await Promise.all([
       fetch("http://localhost:4000/?search[$ne]=null", {
         signal: AbortSignal.timeout(5000),
       }),
       fetch("http://localhost:4000/?search=title", {
         signal: AbortSignal.timeout(5000),
       }),
-    ])
-      .then((results) => {
-        const [noSQLInjection, normalSearch] = results;
-        t.equal(noSQLInjection.status, 200);
-        t.equal(normalSearch.status, 200);
-        t.match(stdout, /Starting agent/);
-        t.notMatch(stderr, /Aikido guard has blocked a NoSQL injection/);
-      })
-      .catch((error) => {
-        t.fail(error.message);
-      })
-      .finally(() => {
-        server.kill();
-        t.end();
-      });
-  }, 1000);
+    ]);
+
+    t.equal(noSQLInjection.status, 200);
+    t.equal(normalSearch.status, 200);
+    t.match(stdout, /Starting agent/);
+    t.notMatch(stderr, /Aikido guard has blocked a NoSQL injection/);
+  } catch (error) {
+    t.fail(error.message);
+  } finally {
+    await new Promise((resolve) => {
+      server.kill();
+      server.on("close", resolve);
+    });
+  }
 });