import unittest
from wif2ansible.models import Server, Flow
from wif2ansible.inventory import generate_inventory

class TestInventoryKeys(unittest.TestCase):
    def test_inventory_keys_are_hostnames(self):
        # Create a server with Ref, Hostname, IP
        s1 = Server(reference="SERVER_REF_01", hostname="server01", ip_address="192.168.1.10", platform="windows")
        
        # Create a flow matching this server
        f1 = Flow(flow_id="1", source_ip="192.168.1.10", destination_ip="10.0.0.1", ports=[80])
        
        servers = {"SERVER_REF_01": s1}
        flows = [f1]
        
        inventory = generate_inventory(servers, flows)
        
        # Verify stricture
        hosts = inventory['all']['hosts']
        
        # Key should be REFERENCE "SERVER_REF_01" (or hostname/ip fallback)
        self.assertIn("SERVER_REF_01", hosts)
        self.assertNotIn("192.168.1.10", hosts)
        
        # Check variables
        host_vars = hosts["SERVER_REF_01"]
        self.assertEqual(host_vars['ansible_host'], "192.168.1.10")
        self.assertEqual(host_vars['ansible_connection'], "winrm")

if __name__ == '__main__':
    unittest.main()