From 12f55ad7212eee99c2361438a4eae02a85cf3d69 Mon Sep 17 00:00:00 2001 From: Abel Luck Date: Mon, 6 Nov 2023 09:01:18 +0100 Subject: [PATCH] add tests --- tailscalesd/main.py | 24 ++++----- tests/test_main.py | 122 ++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 133 insertions(+), 13 deletions(-) create mode 100644 tests/test_main.py diff --git a/tailscalesd/main.py b/tailscalesd/main.py index 842c2a2..334e81e 100644 --- a/tailscalesd/main.py +++ b/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, - "labels": tailscale_labels(device) + "targets": [f"{ipv4}:{port}"], + "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) - matrix_targets = await matrix_sd(devices) + device_targets = plain_devices_sd(settings.tailnet, devices) + matrix_targets = await matrix_sd(settings.tailnet, devices) CACHE_SD = matrix_targets + device_targets await asyncio.sleep(settings.interval) except Exception as e: diff --git a/tests/test_main.py b/tests/test_main.py new file mode 100644 index 0000000..dfb5a37 --- /dev/null +++ b/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) == []