From 1de3a71593d64a2188a28106b2c871001a0af8c6 Mon Sep 17 00:00:00 2001
From: Szymon Marczak <36894700+szmarczak@users.noreply.github.com>
Date: Mon, 29 Aug 2022 12:26:29 +0200
Subject: [PATCH] test: fix DNS cancel tests

---
 .../test-dns-channel-cancel-promise.js        | 22 +++++++------------
 test/parallel/test-dns-channel-cancel.js      | 19 ++++++----------
 2 files changed, 15 insertions(+), 26 deletions(-)

diff --git a/test/parallel/test-dns-channel-cancel-promise.js b/test/parallel/test-dns-channel-cancel-promise.js
index 382ac3dd508d68..6dee3e6a778687 100644
--- a/test/parallel/test-dns-channel-cancel-promise.js
+++ b/test/parallel/test-dns-channel-cancel-promise.js
@@ -7,23 +7,15 @@ const dgram = require('dgram');
 const server = dgram.createSocket('udp4');
 const resolver = new dnsPromises.Resolver();
 
-const addMessageListener = () => {
-  server.removeAllListeners('message');
-
-  server.once('message', () => {
-    server.once('message', common.mustNotCall);
-
-    resolver.cancel();
-  });
-};
-
 server.bind(0, common.mustCall(async () => {
   resolver.setServers([`127.0.0.1:${server.address().port}`]);
 
-  addMessageListener();
-
   // Single promise
   {
+    server.once('message', () => {
+      resolver.cancel();
+    });
+
     const hostname = 'example0.org';
 
     await assert.rejects(
@@ -36,10 +28,12 @@ server.bind(0, common.mustCall(async () => {
     );
   }
 
-  addMessageListener();
-
   // Multiple promises
   {
+    server.once('message', () => {
+      resolver.cancel();
+    });
+
     const assertions = [];
     const assertionCount = 10;
 
diff --git a/test/parallel/test-dns-channel-cancel.js b/test/parallel/test-dns-channel-cancel.js
index f92fb2e30a4d12..405b31e4cc1913 100644
--- a/test/parallel/test-dns-channel-cancel.js
+++ b/test/parallel/test-dns-channel-cancel.js
@@ -10,16 +10,6 @@ const resolver = new Resolver();
 const desiredQueries = 11;
 let finishedQueries = 0;
 
-const addMessageListener = () => {
-  server.removeAllListeners('message');
-
-  server.once('message', () => {
-    server.once('message', common.mustNotCall);
-
-    resolver.cancel();
-  });
-};
-
 server.bind(0, common.mustCall(async () => {
   resolver.setServers([`127.0.0.1:${server.address().port}`]);
 
@@ -37,7 +27,9 @@ server.bind(0, common.mustCall(async () => {
   const next = (...args) => {
     callback(...args);
 
-    addMessageListener();
+    server.once('message', () => {
+      resolver.cancel();
+    });
 
     // Multiple queries
     for (let i = 1; i < desiredQueries; i++) {
@@ -45,7 +37,10 @@ server.bind(0, common.mustCall(async () => {
     }
   };
 
+  server.once('message', () => {
+    resolver.cancel();
+  });
+
   // Single query
-  addMessageListener();
   resolver.resolve4('example0.org', next);
 }));