test: fix DNS cancel tests

PR-URL: #44432
Reviewed-By: Joyee Cheung <joyeec9h3@gmail.com>
Reviewed-By: Luigi Pinca <luigipinca@gmail.com>

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;
 

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,15 +27,20 @@ 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++) {
       resolver.resolve4(`example${i}.org`, callback);
     }
   };
 
+  server.once('message', () => {
+    resolver.cancel();
+  });
+
   // Single query
-  addMessageListener();
   resolver.resolve4('example0.org', next);
 }));