more coverage for wire

1 files changed, 33 insertions, 1 deletions

diff --git a/tests/test_wire.py b/tests/test_wire.py
index ac4401c..2cfaa1b 100644
--- a/tests/test_wire.py
+++ b/tests/test_wire.py
@@ -51,5 +51,37 @@ class BinaryTestCase(unittest.TestCase):
         self.assertEqual(p.get_name(), expected1)
         self.assertEqual(p.get_name(), expected2)
         self.assertEqual(p.remaining(), 0)
-        # verify the unseek()
+        # verify the restore_furthest()
         self.assertEqual(p.current, len(wire))
+
+    def test_seek(self):
+        wire = b'\x09dnspython\x03org\x00'
+        p = dns.wire.Parser(wire)
+        p.seek(10)
+        self.assertEqual(p.get_uint8(), 3)
+        # seeking to the end index is OK
+        p.seek(len(wire))
+        self.assertEqual(p.current, p.end)
+        with self.assertRaises(dns.exception.FormError):
+            # but reading there will not succeed
+            p.get_uint8()
+        with self.assertRaises(dns.exception.FormError):
+            p.seek(-1)
+        with self.assertRaises(dns.exception.FormError):
+            p.seek(len(wire) + 1)
+
+    def test_not_reading_everything_in_restriction(self):
+        wire = bytes.fromhex('0102010203040102')
+        p = dns.wire.Parser(wire)
+        with self.assertRaises(dns.exception.FormError):
+            with p.restrict_to(5):
+                v = p.get_uint8()
+                self.assertEqual(v, 1)
+                # don't read the other 4 bytes
+
+    def test_restriction_does_not_mask_exception(self):
+        wire = bytes.fromhex('0102010203040102')
+        p = dns.wire.Parser(wire)
+        with self.assertRaises(NotImplementedError):
+            with p.restrict_to(5):
+                raise NotImplementedError