Validate inventory hostnames via DNS resolution

test_inventory_keys.py

@@ -1,9 +1,14 @@
 import unittest
+import unittest.mock
 from wif2ansible.models import Server, Flow
 from wif2ansible.inventory import generate_inventory
 
 class TestInventoryKeys(unittest.TestCase):
-    def test_inventory_keys_are_hostnames(self):
+    @unittest.mock.patch('wif2ansible.inventory.is_valid_hostname')
+    def test_inventory_keys_are_hostnames(self, mock_resolves):
+        # Mock DNS to say server01 exists
+        mock_resolves.return_value = True
+        
         # Create a server with Ref, Hostname, IP
         s1 = Server(reference="SERVER_REF_01", hostname="server01", ip_addresses=["192.168.1.10"], platform="windows")
         
@@ -27,20 +32,28 @@ class TestInventoryKeys(unittest.TestCase):
         self.assertEqual(host_vars['ansible_host'], "192.168.1.10")
         self.assertEqual(host_vars['ansible_connection'], "winrm")
 
-    def test_clean_reference_logic(self):
-        from wif2ansible.parsers import clean_reference
+    @unittest.mock.patch('wif2ansible.inventory.is_valid_hostname')
+    def test_inventory_keys_resolution(self, mock_resolves):
+        # Setup mock: 'bad_name' -> False, 'good_name' -> True
+        def side_effect(name):
+            if name == "bad_name": return False
+            if name == "good_name": return True
+            return False
+        mock_resolves.side_effect = side_effect
         
-        # Test cases
-        self.assertEqual(clean_reference("SRV123 MyServer"), "MyServer")
-        self.assertEqual(clean_reference("SVR999 AnotherServer"), "AnotherServer")
-        self.assertEqual(clean_reference("srv001 lowercase"), "lowercase")
-        self.assertEqual(clean_reference("SvR555 MixedCase"), "MixedCase")
-        self.assertEqual(clean_reference("JustName"), "JustName")
-        self.assertEqual(clean_reference("SRV123"), "") # Should be empty? or handle? 
-        # If it's just SRV123, strip returns empty. 
-        # User said "never include these in output". 
-        # If the server is ONLY named SRV123, what then? 
-        # Assuming there is usually a name. 
+        # Server with a BAD hostname but a GOOD reference (simulated)
+        # Actually logic is candidates: [hostname, cleaned_ref, rev_dns]
+        # Let's say hostname is "bad_name" and cleaned ref is "good_name"
+        s1 = Server(reference="SRV01 good_name", hostname="bad_name", ip_addresses=["10.10.10.10"])
         
+        f1 = Flow(flow_id="1", source_ip="10.10.10.10", destination_ip="1.1.1.1", ports=[80])
+        
+        inventory = generate_inventory({"k":s1}, [f1])
+        hosts = inventory['all']['hosts']
+        
+        # It should have skipped "bad_name" and picked "good_name" (from cleaned ref)
+        self.assertIn("good_name", hosts)
+        self.assertNotIn("bad_name", hosts)
+
 if __name__ == '__main__':
     unittest.main()