tests/integration: support 'REPLY' step type

tests/pydnstest/scenario.py

@@ -129,7 +129,7 @@ class Entry:
         rdtype = args.pop(0)
         rr = dns.rrset.from_text(owner, ttl, rdclass, rdtype)
         if len(args) > 0:
-            rd = dns.rdata.from_text(rr.rdclass, rr.rdtype, ' '.join(args), origin = dns.name.from_text(self.origin), relativize = False)
+            rd = dns.rdata.from_text(rr.rdclass, rr.rdtype, ' '.join(args), origin=dns.name.from_text(self.origin), relativize=False)
             rr.add(rd)
         return rr
 
@@ -219,6 +219,8 @@ class Step:
             return self.__check_answer(ctx)
         elif self.type == 'TIME_PASSES':
             return self.__time_passes(ctx)
+        elif self.type == 'REPLY':
+            pass
         else:
             raise Exception('step %s unsupported' % self.type)
 
@@ -246,7 +248,6 @@ class Step:
         ctx.scenario.time = int(self.args[1])
         ctx.set_time(ctx.scenario.time)
 
-
 class Scenario:
     def __init__(self, info):
         """ Initialize scenario with description. """
@@ -268,6 +269,17 @@ class Scenario:
         for rng in self.ranges:
             if rng.eligible(step_id, address):
                 return rng.reply(query)
+        # Find any prescripted one-shot replies
+        for step in self.steps:
+            if step.id <= step_id or step.type != 'REPLY':
+                continue
+            try:
+                candidate = step.data[0]
+                candidate.match(query)
+                step.data.remove(candidate)
+                return candidate.adjust_reply(query)
+            except:
+                pass
 
     def play(self, ctx):
         """ Play given scenario. """
@@ -281,5 +293,3 @@ class Scenario:
                 step.play(ctx)
         except Exception as e:
             raise Exception('step #%d %s' % (step.id, str(e)))
-
-

tests/test_integration.py

 #!/usr/bin/env python
-import sys, os, fileinput
+import sys
+import os
+import fileinput
 from pydnstest import scenario, testserver, test
 import _test_integration as mock_ctx
 
@@ -17,7 +19,7 @@ def get_next(file_in):
             return False
         for csep in (';', '#'):
             if csep in line:
-                line = line[0 : line.index(csep)]
+                line = line[0:line.index(csep)]
         tokens = ' '.join(line.strip().split()).split()
         if len(tokens) == 0:
             continue  # Skip empty lines
@@ -88,6 +90,7 @@ def parse_scenario(op, args, file_in):
             out.steps.append(parse_step(op, args, file_in))
     return out
 
+
 def parse_file(file_in):
     """ Parse scenario from a file. """
     try:
@@ -146,6 +149,7 @@ def play_object(path):
         server.stop()
         mock_ctx.deinit()
 
+
 def test_platform(*args):
     if sys.platform == 'windows':
         raise Exception('not supported at all on Windows')