From 3c582c883e9d4520d8e42a4541c4e8686a209521 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Fran=C3=A7ois=20De=20Keersmaeker?=
 <francois.dekeersmaeker@uclouvain.be>
Date: Wed, 13 Nov 2024 12:22:07 +0100
Subject: [PATCH] Added unit test with a sample DNS cache file

---
 .gitignore                                    |  1 +
 dns_unbound_cache_reader/__init__.py          |  7 +++-
 .../dns_unbound_cache_reader.py               |  3 +-
 test/.gitignore                               |  2 ++
 test/__init__.py                              |  3 ++
 test/sample_dns_cache.txt                     | 15 ++++++++
 test/test_sample_file.py                      | 35 +++++++++++++++++++
 7 files changed, 63 insertions(+), 3 deletions(-)
 create mode 100644 test/.gitignore
 create mode 100644 test/__init__.py
 create mode 100644 test/sample_dns_cache.txt
 create mode 100644 test/test_sample_file.py

diff --git a/.gitignore b/.gitignore
index 5131e13..8f4bc89 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,4 +1,5 @@
 .vscode
 .venv
 __pycache__
+*.txt
 test.py
diff --git a/dns_unbound_cache_reader/__init__.py b/dns_unbound_cache_reader/__init__.py
index 2281e6c..547d4bd 100644
--- a/dns_unbound_cache_reader/__init__.py
+++ b/dns_unbound_cache_reader/__init__.py
@@ -1 +1,6 @@
-from .dns_unbound_cache_reader import read_dns_cache
+from .dns_unbound_cache_reader import (
+    DnsCacheSection,
+    DnsRtype,
+    DnsTableKeys,
+    read_dns_cache
+)
diff --git a/dns_unbound_cache_reader/dns_unbound_cache_reader.py b/dns_unbound_cache_reader/dns_unbound_cache_reader.py
index aef3c92..2fe8119 100644
--- a/dns_unbound_cache_reader/dns_unbound_cache_reader.py
+++ b/dns_unbound_cache_reader/dns_unbound_cache_reader.py
@@ -22,7 +22,7 @@ to_skip = (
     "EOF"
 )
 # Regex patterns
-pattern_line      = r"^([a-zA-Z0-9.-]+)\s+(\d+)\s+IN\s+([A-Z]+)\s+(.+)$"                    # Generic DNS cache line
+pattern_line      = r"^([a-zA-Z0-9._-]+)\s+(\d+)\s+IN\s+([A-Z]+)\s+(.+)$"                    # Generic DNS cache line
 pattern_ipv4_byte = r"(25[0-5]|2[0-4][0-9]|1[0-9]{2}|[1-9]?[0-9])"                          # Single byte from an IPv4 address
 pattern_ptr       = (pattern_ipv4_byte + r"\.") * 3 + pattern_ipv4_byte + r".in-addr.arpa"  # Reverse DNS lookup qname
 pattern_srv       = r"^(\d+)\s+(\d+)\s+(\d+)\s+([a-zA-Z0-9.-]+)$"                           # SRV record target
@@ -41,7 +41,6 @@ class DnsRtype(IntEnum):
     Enum class for the DNS resource record types.
     """
     A    = 1   # IPv4 address
-    NS   = 2   # Name server
     PTR  = 12  # Domain name pointer
     AAAA = 28  # IPv6 address
     SRV  = 33  # Service locator
diff --git a/test/.gitignore b/test/.gitignore
new file mode 100644
index 0000000..31706f2
--- /dev/null
+++ b/test/.gitignore
@@ -0,0 +1,2 @@
+# Exception: sample text files for unit testing
+!*.txt
diff --git a/test/__init__.py b/test/__init__.py
new file mode 100644
index 0000000..92283cc
--- /dev/null
+++ b/test/__init__.py
@@ -0,0 +1,3 @@
+"""
+Unit tests for the package `dns_unbound_cache_reader`.
+"""
diff --git a/test/sample_dns_cache.txt b/test/sample_dns_cache.txt
new file mode 100644
index 0000000..a78f062
--- /dev/null
+++ b/test/sample_dns_cache.txt
@@ -0,0 +1,15 @@
+; unbound cache dump
+START_RRSET_CACHE
+example.com.   300    IN    A    93.184.216.34
+example.com.   300    IN    AAAA 2606:2800:220:1:248:1893:25c8:1946
+example.org.   600    IN    MX   10 mail.example.org.
+example.org.   600    IN    NS   ns1.example.org.
+example.org.   600    IN    NS   ns2.example.org.
+ns1.example.org.  3600  IN    A    192.0.2.1
+ns2.example.org.  3600  IN    A    198.51.100.1
+1.2.0.192.in-addr.arpa. 3600    IN  PTR example.com
+2.2.0.192.in-addr.arpa. 3600    IN  PTR example.com
+_tcp_.matter.example.com.   3600    IN  SRV 10 60 5000 server1.example.com
+_tcp_.matter.example.com.   3600    IN  SRV 20 60 5000 server2.example.com
+END_RRSET_CACHE
+EOF
diff --git a/test/test_sample_file.py b/test/test_sample_file.py
new file mode 100644
index 0000000..b89fde7
--- /dev/null
+++ b/test/test_sample_file.py
@@ -0,0 +1,35 @@
+# Libraries
+import os
+# Package under test
+import dns_unbound_cache_reader as dns_reader
+from dns_unbound_cache_reader import DnsTableKeys
+
+
+# Variables
+self_path = os.path.abspath(__file__)
+self_dir = os.path.dirname(self_path)
+sample_cache_file = os.path.join(self_dir, "sample_dns_cache.txt")
+
+
+### TEST FUNCTIONS ###
+
+def test_read_sample_cache_file() -> None:
+    """
+    Test reading a sample cache file.
+    """
+    dns_table = dns_reader.read_dns_cache(file=sample_cache_file)
+    assert DnsTableKeys.IP.name in dns_table
+    assert DnsTableKeys.SERVICE.name in dns_table
+
+    dns_table_ip = dns_table[DnsTableKeys.IP.name]
+    assert len(dns_table_ip) == 5
+    assert dns_table_ip["93.184.216.34"] == "example.com"
+    assert dns_table_ip["2606:2800:220:1:248:1893:25c8:1946"] == "example.com"
+    assert dns_table_ip["192.0.2.1"] == "ns1.example.org"
+    assert dns_table_ip["198.51.100.1"] == "ns2.example.org"
+    assert dns_table_ip["192.0.2.2"] == "example.com"
+    
+    dns_table_service = dns_table[DnsTableKeys.SERVICE.name]
+    assert len(dns_table_service) == 2
+    assert dns_table_service["server1.example.com"] == "_tcp_.matter.example.com"
+    assert dns_table_service["server2.example.com"] == "_tcp_.matter.example.com"
-- 
GitLab