add tests

tailscalesd/main.py

@@ -70,7 +70,7 @@ def group_by_type(input_list):
     return result
 
 
-def tailscale_labels(device) -> Dict[str, str]:
+def tailscale_labels(tailnet, device) -> Dict[str, str]:
     return {
         "__meta_tailscale_device_client_version": device["clientVersion"],
         "__meta_tailscale_device_hostname": device["hostname"],
@@ -78,7 +78,7 @@ def tailscale_labels(device) -> Dict[str, str]:
         "__meta_tailscale_device_id": device["id"],
         "__meta_tailscale_device_name": device["name"],
         "__meta_tailscale_device_os": device["os"],
-        "__meta_tailscale_tailnet": settings.tailnet,
+        "__meta_tailscale_tailnet": tailnet,
     }
 
 
@@ -90,13 +90,12 @@ async def matrix_node_sd(device) -> Dict:
         return group_by_type(data)
 
 
-def matrix_workers_to_sd(device, workers) -> List:
+def matrix_workers_to_sd(tailnet, device, workers) -> List:
     if len(workers) == 0:
         return []
     ipv4 = ipv4_only(device["addresses"])[0]
     target_groups = []
     for worker_type, workers in workers.items():
-        targets = []
         for worker in workers:
             port = worker["metrics_port"]
             worker_name = worker.get("name", "WORKER_NO_NAME")
@@ -105,11 +104,10 @@ def matrix_workers_to_sd(device, workers) -> List:
                     f"Error parsing worker {worker_name} on host={device['hostname']}. Port is invalid port={port}"
                 )
                 continue
-            targets.append(f"{ipv4}:{port}")
             target_groups.append(
                 {
-                    "targets": targets,
+                    "targets": [f"{ipv4}:{port}"],
-                    "labels": tailscale_labels(device)
+                    "labels": tailscale_labels(tailnet, device)
                     | {
                         "__meta_matrix_worker_type": worker_type,
                         "__meta_matrix_worker_name": worker_name,
@@ -119,7 +117,7 @@ def matrix_workers_to_sd(device, workers) -> List:
     return target_groups
 
 
-async def matrix_sd(devices) -> List:
+async def matrix_sd(tailnet, devices) -> List:
     sd = []
     for device in devices:
         if "tag:matrix" not in device["tags"]:
@@ -132,17 +130,17 @@ async def matrix_sd(devices) -> List:
                 exc_info=e,
             )
             workers = {}
-        targets = matrix_workers_to_sd(device, workers)
+        targets = matrix_workers_to_sd(tailnet, device, workers)
         if targets:
             sd.append(targets)
     return []
 
 
-def plain_devices_sd(devices) -> List:
+def plain_devices_sd(tailnet, devices) -> List:
     sd = []
     for device in devices:
         targets = ipv4_only(device["addresses"])
-        labels = tailscale_labels(device)
+        labels = tailscale_labels(tailnet, device)
         sd.append({"labels": labels, "targets": targets})
 
     return sd
@@ -153,8 +151,8 @@ async def poll_sd():
     while True:
         try:
             devices = await tailscale_devices()
-            device_targets = plain_devices_sd(devices)
+            device_targets = plain_devices_sd(settings.tailnet, devices)
-            matrix_targets = await matrix_sd(devices)
+            matrix_targets = await matrix_sd(settings.tailnet, devices)
             CACHE_SD = matrix_targets + device_targets
             await asyncio.sleep(settings.interval)
         except Exception as e:

tests/test_main.py

@@ -0,0 +1,122 @@
+from tailscalesd.main import group_by_type, matrix_workers_to_sd
+
+
+def test_group_by_type_valid_input():
+    input_list = [
+        {"type": "A", "value": 1},
+        {"type": "B", "value": 2},
+        {"type": "A", "value": 3},
+        {"type": "C", "value": 5},
+    ]
+    expected_output = {
+        "A": [{"type": "A", "value": 1}, {"type": "A", "value": 3}],
+        "B": [{"type": "B", "value": 2}],
+        "C": [{"type": "C", "value": 5}],
+    }
+    assert group_by_type(input_list) == expected_output
+
+
+def test_group_by_type_missing_type_key():
+    input_list = [
+        {"value": 1},
+        {"type": "B", "value": 2},
+        {"value": 3},
+        {"type": "C", "value": 5},
+    ]
+    expected_output = {
+        "B": [{"type": "B", "value": 2}],
+        "C": [{"type": "C", "value": 5}],
+    }
+    assert group_by_type(input_list) == expected_output
+
+
+def test_group_by_type_empty_input():
+    input_list = []
+    expected_output = {}
+    assert group_by_type(input_list) == expected_output
+
+
+def test_matrix_workers_to_sd_valid_input():
+    device = {
+        "addresses": ["10.0.0.1", "192.168.1.1"],
+        "hostname": "test-host",
+        "clientVersion": "1",
+        "tags": ["tag:matrix"],
+        "id": "id",
+        "name": "name",
+        "os": "linux",
+        "authorized": True,
+        "tailnet": "testnet",
+    }
+    workers = {
+        "type1": [
+            {"metrics_port": 1234, "name": "worker1"},
+            {"metrics_port": 1235, "name": "worker2"},
+        ],
+        "type2": [
+            {"metrics_port": 1236, "name": "worker3"},
+        ],
+    }
+    ts_labels = {
+        "__meta_tailscale_device_client_version": "1",
+        "__meta_tailscale_device_hostname": "test-host",
+        "__meta_tailscale_device_authorized": "true",
+        "__meta_tailscale_device_id": "id",
+        "__meta_tailscale_device_name": "name",
+        "__meta_tailscale_device_os": "linux",
+        "__meta_tailscale_tailnet": "testnet",
+    }
+    expected_output = [
+        {
+            "targets": ["10.0.0.1:1234"],
+            "labels": {
+                # Add expected tailscale_labels output here
+                "__meta_matrix_worker_type": "type1",
+                "__meta_matrix_worker_name": "worker1",
+            }
+            | ts_labels,
+        },
+        {
+            "targets": ["10.0.0.1:1235"],
+            "labels": {
+                # Add expected tailscale_labels output here
+                "__meta_matrix_worker_type": "type1",
+                "__meta_matrix_worker_name": "worker2",
+            }
+            | ts_labels,
+        },
+        {
+            "targets": ["10.0.0.1:1236"],
+            "labels": {
+                # Add expected tailscale_labels output here
+                "__meta_matrix_worker_type": "type2",
+                "__meta_matrix_worker_name": "worker3",
+            }
+            | ts_labels,
+        },
+    ]
+    actual_output = matrix_workers_to_sd("testnet", device, workers)
+    assert actual_output == expected_output
+
+
+def test_matrix_workers_to_sd_empty_workers():
+    device = {
+        "addresses": ["10.0.0.1", "192.168.1.1"],
+        "hostname": "test-host",
+    }
+    workers = {}
+    assert matrix_workers_to_sd("testnet", device, workers) == []
+
+
+def test_matrix_workers_to_sd_invalid_port():
+    device = {
+        "addresses": ["10.0.0.1", "192.168.1.1"],
+        "hostname": "test-host",
+    }
+    workers = {
+        "type1": [
+            {"metrics_port": None, "name": "worker1"},
+        ],
+    }
+    # Assuming log.error does not raise any exceptions
+    assert matrix_workers_to_sd("testnet", device, workers) == []