Merge pull request #8713 from rgacogne/auth-strict-caches-size

[thirdparty/pdns.git] / regression-tests.dnsdist / dnsdisttests.py

diff --git a/regression-tests.dnsdist/dnsdisttests.py b/regression-tests.dnsdist/dnsdisttests.py
index 792ceefcc63bc94542a218b30a11e00a619c8069..01eb5332edab98a2dc12818aba7df39f7b7a6ce4 100644 (file)
--- a/regression-tests.dnsdist/dnsdisttests.py
+++ b/regression-tests.dnsdist/dnsdisttests.py
@@ -16,6 +16,8 @@ import dns.message
 import libnacl
 import libnacl.utils
 
+from eqdnsmessage import AssertEqualDNSMessageMixin
+
 # Python2/3 compatibility hacks
 try:
   from queue import Queue
@@ -28,7 +30,7 @@ except NameError:
   pass
 
 
-class DNSDistTest(unittest.TestCase):
+class DNSDistTest(AssertEqualDNSMessageMixin, unittest.TestCase):
     """
     Set up a dnsdist instance and responder threads.
     Queries sent to dnsdist are relayed to the responder threads,
@@ -55,6 +57,7 @@ class DNSDistTest(unittest.TestCase):
     _healthCheckName = 'a.root-servers.net.'
     _healthCheckCounter = 0
     _answerUnexpected = True
+    _checkConfigExpectedOutput = None
 
     @classmethod
     def startResponders(cls):
@@ -89,7 +92,10 @@ class DNSDistTest(unittest.TestCase):
             output = subprocess.check_output(testcmd, stderr=subprocess.STDOUT, close_fds=True)
         except subprocess.CalledProcessError as exc:
             raise AssertionError('dnsdist --check-config failed (%d): %s' % (exc.returncode, exc.output))
-        expectedOutput = ('Configuration \'%s\' OK!\n' % (confFile)).encode()
+        if cls._checkConfigExpectedOutput is not None:
+          expectedOutput = cls._checkConfigExpectedOutput
+        else:
+          expectedOutput = ('Configuration \'%s\' OK!\n' % (confFile)).encode()
         if output != expectedOutput:
             raise AssertionError('dnsdist --check-config failed: %s' % output)
 
@@ -198,6 +204,7 @@ class DNSDistTest(unittest.TestCase):
                 request = dns.message.from_wire(data, ignore_trailing=True)
                 forceRcode = trailingDataResponse
 
+            wire = None
             if callback:
               wire = callback(request)
             else:
@@ -205,6 +212,9 @@ class DNSDistTest(unittest.TestCase):
               if response:
                 wire = response.to_wire()
 
+            if not wire:
+              continue
+
             sock.settimeout(2.0)
             sock.sendto(wire, addr)
             sock.settimeout(None)
@@ -450,6 +460,8 @@ class DNSDistTest(unittest.TestCase):
         while not self._fromResponderQueue.empty():
             self._fromResponderQueue.get(False)
 
+        super(DNSDistTest, self).setUp()
+
     @classmethod
     def clearToResponderQueue(cls):
         while not cls._toResponderQueue.empty():
@@ -573,3 +585,4 @@ class DNSDistTest(unittest.TestCase):
 
     def checkResponseNoEDNS(self, expected, received):
         self.checkMessageNoEDNS(expected, received)
+