func-tests: change answer_count to more general count

tests-extra/tests/dnssec/dnskey_timestamps/test.py

@@ -33,8 +33,8 @@ def check_zone(server, expect_dnskey, expect_rrsig):
     dnskeys = server.dig("example.com", "DNSKEY")
     soa = server.dig("example.com", "SOA", dnssec=True)
 
-    found_dnskeys = dnskeys.answer_count("DNSKEY")
-    found_rrsigs = soa.answer_count("RRSIG")
+    found_dnskeys = dnskeys.count("DNSKEY")
+    found_rrsigs = soa.count("RRSIG")
 
     expect_dnskeys = 2 if expect_dnskey else 1
     expect_rrsigs = 2 if expect_rrsig else 1

tests-extra/tests/dnssec/none_to_nsec/test.py

@@ -38,11 +38,11 @@ t.sleep(1)
 
 # Check absence of NSEC3PARAM record.
 resp = master.dig(zone, "NSEC3PARAM", dnssec=True)
-compare(resp.answer_count(), 0, "NSEC3PARAM count")
+compare(resp.count(), 0, "NSEC3PARAM count")
 
 # Check presence of DNSKEYs.
 resp = master.dig(zone, "DNSKEY", dnssec=True)
-compare(resp.answer_count(), 4, "DNSKEY count")
+compare(resp.count(), 4, "DNSKEY count")
 
 # Check NSEC presence.
 master.check_nsec(zone)

tests-extra/tests/dnssec/none_to_nsec3/test.py

@@ -39,11 +39,11 @@ t.sleep(1)
 
 # Check presence of NSEC3PARAM record.
 resp = master.dig(zone, "NSEC3PARAM", dnssec=True)
-compare(resp.answer_count(), 1, "NSEC3PARAM count")
+compare(resp.count(), 1, "NSEC3PARAM count")
 
 # Check presence of DNSKEYs.
 resp = master.dig(zone, "DNSKEY", dnssec=True)
-compare(resp.answer_count(), 2, "DNSKEY count")
+compare(resp.count(), 2, "DNSKEY count")
 
 # Check NSEC presence.
 master.check_nsec(zone, nsec3=True)

tests-extra/tests/dnssec/nsec3_to_nsec/test.py

@@ -39,11 +39,11 @@ t.sleep(1)
 
 # Check presence of NSEC3PARAM record.
 resp = master.dig(zone, "NSEC3PARAM", dnssec=True)
-compare(resp.answer_count(), 1, "NSEC3PARAM count")
+compare(resp.count(), 1, "NSEC3PARAM count")
 
 # Check presence of DNSKEYs.
 resp = master.dig(zone, "DNSKEY", dnssec=True)
-compare(resp.answer_count(), 4, "DNSKEY count")
+compare(resp.count(), 4, "DNSKEY count")
 
 # Check NSEC3 presence.
 master.check_nsec(zone, nsec3=True)
@@ -69,11 +69,11 @@ t.sleep(1)
 
 # Check absence of NSEC3PARAM record.
 resp = master.dig(zone, "NSEC3PARAM", dnssec=True)
-compare(resp.answer_count(), 0, "NSEC3PARAM count")
+compare(resp.count(), 0, "NSEC3PARAM count")
 
 # Check presence of DNSKEYs.
 resp = master.dig(zone, "DNSKEY", dnssec=True)
-compare(resp.answer_count(), 4, "DNSKEY count")
+compare(resp.count(), 4, "DNSKEY count")
 
 # Check NSEC presence.
 master.check_nsec(zone)

tests-extra/tests/dnssec/nsec3_to_nsec_ddns/test.py

@@ -39,11 +39,11 @@ t.sleep(1)
 
 # Check presence of NSEC3PARAM record.
 resp = master.dig(zone, "NSEC3PARAM", dnssec=True)
-compare(resp.answer_count(), 1, "NSEC3PARAM count")
+compare(resp.count(), 1, "NSEC3PARAM count")
 
 # Check presence of DNSKEYs.
 resp = master.dig(zone, "DNSKEY", dnssec=True)
-compare(resp.answer_count(), 4, "DNSKEY count")
+compare(resp.count(), 4, "DNSKEY count")
 
 # Check NSEC3 presence.
 master.check_nsec(zone, nsec3=True)
@@ -68,11 +68,11 @@ t.sleep(1)
 
 # Check absence of NSEC3PARAM record.
 resp = master.dig(zone, "NSEC3PARAM", dnssec=True)
-compare(resp.answer_count(), 0, "NSEC3PARAM count")
+compare(resp.count(), 0, "NSEC3PARAM count")
 
 # Check presence of DNSKEYs.
 resp = master.dig(zone, "DNSKEY", dnssec=True)
-compare(resp.answer_count(), 4, "DNSKEY count")
+compare(resp.count(), 4, "DNSKEY count")
 
 # Check NSEC presence.
 master.check_nsec(zone)

tests-extra/tests/dnssec/nsec_to_nsec3/test.py

@@ -38,11 +38,11 @@ t.sleep(1)
 
 # Check absence of NSEC3PARAM record.
 resp = master.dig(zone, "NSEC3PARAM", dnssec=True)
-compare(resp.answer_count(), 0, "NSEC3PARAM count")
+compare(resp.count(), 0, "NSEC3PARAM count")
 
 # Check presence of DNSKEYs.
 resp = master.dig(zone, "DNSKEY", dnssec=True)
-compare(resp.answer_count(), 4, "DNSKEY count")
+compare(resp.count(), 4, "DNSKEY count")
 
 # Check NSEC presence.
 master.check_nsec(zone)
@@ -68,11 +68,11 @@ t.sleep(1)
 
 # Check presence of NSEC3PARAM record.
 resp = master.dig(zone, "NSEC3PARAM", dnssec=True)
-compare(resp.answer_count(), 1, "NSEC3PARAM count")
+compare(resp.count(), 1, "NSEC3PARAM count")
 
 # Check presence of DNSKEYs.
 resp = master.dig(zone, "DNSKEY", dnssec=True)
-compare(resp.answer_count(), 4, "DNSKEY count")
+compare(resp.count(), 4, "DNSKEY count")
 
 # Check NSEC3 presence.
 master.check_nsec(zone, nsec3=True)

tests-extra/tests/dnssec/nsec_to_nsec3_ddns/test.py

@@ -38,11 +38,11 @@ t.sleep(1)
 
 # Check absence of NSEC3PARAM record.
 resp = master.dig(zone, "NSEC3PARAM", dnssec=True)
-compare(resp.answer_count(), 0, "NSEC3PARAM count")
+compare(resp.count(), 0, "NSEC3PARAM count")
 
 # Check presence of DNSKEYs.
 resp = master.dig(zone, "DNSKEY", dnssec=True)
-compare(resp.answer_count(), 4, "DNSKEY count")
+compare(resp.count(), 4, "DNSKEY count")
 
 # Check NSEC presence.
 master.check_nsec(zone)
@@ -67,11 +67,11 @@ t.sleep(1)
 
 # Check presence of NSEC3PARAM record.
 resp = master.dig(zone, "NSEC3PARAM", dnssec=True)
-compare(resp.answer_count(), 1, "NSEC3PARAM count")
+compare(resp.count(), 1, "NSEC3PARAM count")
 
 # Check presence of DNSKEYs.
 resp = master.dig(zone, "DNSKEY", dnssec=True)
-compare(resp.answer_count(), 4, "DNSKEY count")
+compare(resp.count(), 4, "DNSKEY count")
 
 # Check NSEC3 presence.
 master.check_nsec(zone, nsec3=True)

tests-extra/tools/dnstest/response.py

@@ -159,15 +159,22 @@ class Response(object):
         resp = server.dig(**self.args)
         self.diff(resp, flags, answer, authority, additional)
 
-    def answer_count(self, rtype=None):
-        '''Returns number of records of given type in answer section'''
+    def count(self, rtype=None, section="answer"):
+        '''Returns number of records of given type in specified section'''
 
         if not rtype:
             rtype = self.rtype
         elif type(rtype) is str:
             rtype = dns.rdatatype.from_text(rtype)
 
-        for rrset in self.resp.answer:
+        if not section or section == "answer":
+            sect = self.resp.answer
+        elif section == "additional":
+            sect = self.resp.additional
+        elif section == "authority":
+            sect = self.resp.authority
+
+        for rrset in sect:
             if rrset.rdtype == rtype:
                 return len(rrset)
         else: