[dns_check] Allow querying a custom record type (#2653)

New `record_type` optional parameter in the configuration.

Along with tests: runs of the check in `test_invalid_config` now reload
the config correctly and test that the expected exception is raised.

checks.d/dns_check.py

@@ -51,12 +51,14 @@ def query_dns(self, instance, timeout):
         if nameserver is not None:
             resolver.nameservers = [nameserver]
 
+        record_type = instance.get('record_type', 'A')
+
         status = AgentCheck.CRITICAL
         start_time = time.time()
         try:
-            self.log.debug('Resolving hostname %s...' % hostname)
-            answer = resolver.query(hostname)
-            assert(answer.rrset.items[0].address)
+            self.log.debug('Querying "%s" record for hostname "%s"...' % (record_type, hostname))
+            answer = resolver.query(hostname, rdtype=record_type)
+            assert(answer.rrset.items[0].to_text())
             end_time = time.time()
         except dns.exception.Timeout:
             self.log.error('DNS resolution of %s timed out' % hostname)

conf.d/dns_check.yaml.example

@@ -5,3 +5,7 @@ instances:
     - hostname: www.example.org
       nameserver: 127.0.0.1
       timeout: 8
+
+      # Specify an (optional) `record_type` to customize the record type
+      # queried by the check (default: "A")
+      # record_type: A

tests/checks/mock/test_dns_check.py

@@ -2,6 +2,7 @@
 import mock
 
 # 3p
+from dns.rdatatype import UnknownRdatatype
 from dns.resolver import Resolver, Timeout, NXDOMAIN
 
 # project
@@ -22,19 +23,21 @@ def __init__(self, address):
 
     class MockRrset:
         def __init__(self, address):
-            self.items = [MockDNSAnswer.MockItem()]
+            self.items = [MockDNSAnswer.MockItem(address)]
 
     class MockItem:
-        def __getattr__(self, name):
-            return '127.0.0.1'
-
+        def __init__(self, address):
+            self._address = address
 
-def success_query_mock(d_name):
-    return MockDNSAnswer('127.0.0.1')
+        def to_text(self):
+            return self._address
 
 
-def timeout_query_mock(d_name):
-    raise Timeout()
+def success_query_mock(d_name, rdtype):
+    if rdtype == 'A':
+        return MockDNSAnswer('127.0.0.1')
+    elif rdtype == 'CNAME':
+        return MockDNSAnswer('alias.example.org')
 
 
 class TestDns(AgentCheckTest):
@@ -52,6 +55,18 @@ def test_success(self, mocked_query):
                                 tags=['resolved_hostname:www.example.org', 'nameserver:127.0.0.1'])
         self.coverage_report()
 
+    @mock.patch.object(Resolver, 'query', side_effect=success_query_mock)
+    def test_success_CNAME(self, mocked_query):
+        config = {
+            'instances': [{'hostname': 'www.example.org', 'nameserver': '127.0.0.1', 'record_type': 'CNAME'}]
+        }
+        self.run_check(config)
+        self.assertMetric('dns.response_time', count=1,
+                          tags=['nameserver:127.0.0.1', 'resolved_hostname:www.example.org'])
+        self.assertServiceCheck(SERVICE_CHECK_NAME, status=AgentCheck.OK,
+                                tags=['resolved_hostname:www.example.org', 'nameserver:127.0.0.1'])
+        self.coverage_report()
+
     @mock.patch.object(Resolver, 'query', side_effect=Timeout())
     def test_timeout(self, mocked_query):
         configs = [
@@ -74,12 +89,14 @@ def test_timeout(self, mocked_query):
     def test_invalid_config(self):
         configs = [
             # invalid hostname
-            {'instances': [{'hostname': 'example'}]},
+            ({'instances': [{'hostname': 'example'}]}, NXDOMAIN),
             # invalid nameserver
-            {'instances': [{'hostname': 'www.example.org', 'nameserver': '0.0.0.0'}]}
+            ({'instances': [{'hostname': 'www.example.org', 'nameserver': '0.0.0.0'}]}, Timeout),
+            # invalid record type
+            ({'instances': [{'hostname': 'www.example.org', 'record_type': 'FOO'}]}, UnknownRdatatype),
         ]
-        for config in configs:
-            self.assertRaises(NXDOMAIN, lambda: self.run_check(config))
+        for config, exception_class in configs:
+            self.assertRaises(exception_class, self.run_check, config, force_reload=True)
             self.assertEquals(len(self.metrics), 0)
             self.assertServiceCheck(SERVICE_CHECK_NAME, status=AgentCheck.CRITICAL)
             self.coverage_report()