tests/krb5/raw_testcase.py: add assertElement*()

These helper functions make writing subsequent Kerberos test
clearer.

Signed-off-by: Stefan Metzmacher <metze@samba.org>
Reviewed-by: Andrew Bartlett <abartlet@samba.org>

1 files changed, 54 insertions, 0 deletions

diff --git a/python/samba/tests/krb5/raw_testcase.py b/python/samba/tests/krb5/raw_testcase.py
index 333aab70c8e..eb294a75a95 100644
--- a/python/samba/tests/krb5/raw_testcase.py
+++ b/python/samba/tests/krb5/raw_testcase.py
@@ -605,6 +605,36 @@ class RawKerberosTest(TestCaseInTempDir):
         self.assertIsNotNone(value)
         return
 
+    def getElementValue(self, obj, elem):
+        v = None
+        try:
+            v = obj[elem]
+        except KeyError:
+            pass
+        return v
+
+    def assertElementMissing(self, obj, elem):
+        v = self.getElementValue(obj, elem)
+        self.assertIsNone(v)
+        return
+
+    def assertElementPresent(self, obj, elem):
+        v = self.getElementValue(obj, elem)
+        self.assertIsNotNone(v)
+        return
+
+    def assertElementEqual(self, obj, elem, value):
+        v = self.getElementValue(obj, elem)
+        self.assertIsNotNone(v)
+        self.assertEqual(v, value)
+        return
+
+    def assertElementEqualUTF8(self, obj, elem, value):
+        v = self.getElementValue(obj, elem)
+        self.assertIsNotNone(v)
+        self.assertEqual(v, bytes(value, 'utf8'))
+        return
+
     def assertPrincipalEqual(self, princ1, princ2):
         self.assertEqual(princ1['name-type'], princ2['name-type'])
         self.assertEqual(
@@ -618,6 +648,30 @@ class RawKerberosTest(TestCaseInTempDir):
                 msg="princ1=%s != princ2=%s" % (princ1, princ2))
         return
 
+    def assertElementEqualPrincipal(self, obj, elem, value):
+        v = self.getElementValue(obj, elem)
+        self.assertIsNotNone(v)
+        v = pyasn1_native_decode(v, asn1Spec=krb5_asn1.PrincipalName())
+        self.assertPrincipalEqual(v, value)
+        return
+
+    def assertElementKVNO(self, obj, elem, value):
+        v = self.getElementValue(obj, elem)
+        if value == "autodetect":
+            value = v
+        if value is not None:
+            self.assertIsNotNone(v)
+            # The value on the wire should never be 0
+            self.assertNotEqual(v, 0)
+            # value == 0 means we don't know the kvno
+            # but enforce at any value != 0 is present
+            value = int(value)
+            if value != 0:
+                self.assertEqual(v, value)
+        else:
+            self.assertIsNone(v)
+        return
+
     def get_KerberosTimeWithUsec(self, epoch=None, offset=None):
         if epoch is None:
             epoch = time.time()