Add ping check logic

Change-Id: I26ac9205f14d7fb053f6ed4da640664dac4621fd

Author: benfogel

app/check_ping.py

@@ -0,0 +1,162 @@
+import logging
+import re
+
+from kubernetes import client
+from pydantic import BaseModel, Field
+
+from distributed_runner import DistributedRunner
+
+log = logging.getLogger("check.ping")
+
+default_image = "alpine:3.22"
+
+
+class CheckPingParameters(BaseModel):
+    namespace: str = "cluster-health"
+    image: str = default_image
+    secondary_network_config: str = Field(None, alias="secondaryNetworkConfig")
+    timeout_seconds: int = Field(300, alias="timeoutSeconds")
+    count: int = 1
+    avg_rtt_ms_threshold: float = Field(None, alias="avgRttMsThreshold")
+    max_rtt_ms_threshold: float = Field(None, alias="maxRttMsThreshold")
+    ping_target: str = Field(alias="pingTarget")
+
+
+class CheckPing:
+    """
+    Runs a network ping test on each node in the cluster.
+    The check creates a Pod that runs on each node one at a time. Each pod
+    executes 'ping', and the results are read from the pod logs.
+    The packet loss and round-trip times are then checked. The Pod is
+    cleaned up after the check.
+    """
+
+    def __init__(self, parameters: dict) -> None:
+        params = CheckPingParameters(**parameters)
+        self.namespace = params.namespace
+        self.secondary_network_config = params.secondary_network_config
+        self.image = params.image
+        self.count = params.count
+        self.timeout_seconds = params.timeout_seconds
+        self.avg_rtt_ms_threshold = params.avg_rtt_ms_threshold
+        self.max_rtt_ms_threshold = params.max_rtt_ms_threshold
+        self.ping_target = params.ping_target
+        self.k8s_core_v1 = client.CoreV1Api()
+        self.k8s_apps_v1 = client.AppsV1Api()
+
+    def parse_pod_logs_from_ping_test_func(self) -> callable:
+        """
+        Creates a parsing function that reads log output from each pod and assesses health check pass/fail.
+        """
+        count = self.count
+        avg_rtt_ms_threshold = self.avg_rtt_ms_threshold
+        max_rtt_ms_threshold = self.max_rtt_ms_threshold
+
+        def parse_pod_logs_from_ping_test(logs: str, node_name: str) -> bool:
+            # Check for 0% packet loss
+            packet_loss_match = re.search(r"(\d+)% packet loss", logs)
+            if not packet_loss_match:
+                log.error(
+                    f"Could not parse packet loss from ping results on node {node_name}. Logs:\n{logs}"
+                )
+                return False
+
+            packet_loss = int(packet_loss_match.group(1))
+            if packet_loss > 0:
+                log.error(f"Node {node_name}: {packet_loss}% packet loss detected.")
+                return False
+
+            # Check if the correct number of packets were transmitted and received
+            transmitted_match = re.search(r"(\d+) packets transmitted", logs)
+            if not transmitted_match:
+                log.error(
+                    f"Could not parse transmitted packets on node {node_name}. Logs:\n{logs}"
+                )
+                return False
+
+            transmitted = int(transmitted_match.group(1))
+            if transmitted != count:
+                log.error(
+                    f"Node {node_name}: Expected {count} packets transmitted, but got {transmitted}."
+                )
+                return False
+
+            # If RTT checks are configured, parse and validate them.
+            if avg_rtt_ms_threshold is not None or max_rtt_ms_threshold is not None:
+                # Regex for iputils-ping: rtt min/avg/max/mdev = 0.052/0.052/0.052/0.000 ms
+                rtt_match = re.search(
+                    r"rtt min/avg/max/mdev = [\d\.]+/([\d\.]+)/([\d\.]+)/[\d\.]+ ms", logs
+                )
+                if not rtt_match:
+                    # Fallback for busybox ping: round-trip min/avg/max = 0.052/0.052/0.052 ms
+                    rtt_match = re.search(
+                        r"round-trip min/avg/max = [\d\.]+/([\d\.]+)/([\d\.]+) ms", logs
+                    )
+
+                if not rtt_match:
+                    log.error(
+                        f"Could not parse RTT from ping results on node {node_name} but RTT checks were requested. Logs:\n{logs}"
+                    )
+                    return False
+
+                rtt_avg = float(rtt_match.group(1))
+                rtt_max = float(rtt_match.group(2))
+                log.info(
+                    f"Node {node_name}: Ping RTT avg={rtt_avg:.2f}ms, max={rtt_max:.2f}ms"
+                )
+
+                if avg_rtt_ms_threshold is not None and rtt_avg > avg_rtt_ms_threshold:
+                    log.error(
+                        f"Node {node_name} average RTT {rtt_avg:.2f}ms is above maximum of {avg_rtt_ms_threshold}ms."
+                    )
+                    return False
+
+                if max_rtt_ms_threshold is not None and rtt_max > max_rtt_ms_threshold:
+                    log.error(
+                        f"Node {node_name} max RTT {rtt_max:.2f}ms is above maximum of {max_rtt_ms_threshold}ms."
+                    )
+                    return False
+
+            log.info(f"Node {node_name}: Ping test successful with 0% packet loss.")
+            return True
+
+        return parse_pod_logs_from_ping_test
+
+    def is_healthy(self):
+        command_args = ["ping", "-c", str(self.count), self.ping_target]
+
+        container = client.V1Container(
+            name="ping-container",
+            image=self.image,
+            command=command_args,
+        )
+
+        pod_spec = client.V1PodSpec(
+            containers=[container],
+            restart_policy="Never",
+            termination_grace_period_seconds=5,
+        )
+
+        metadata = client.V1ObjectMeta(
+            labels={"app": "ping-test"},
+        )
+
+        if self.secondary_network_config:
+            metadata.annotations = {
+                "networking.gke.io/interfaces": self.secondary_network_config
+            }
+
+        template = client.V1PodTemplateSpec(
+            metadata=metadata,
+            spec=pod_spec,
+        )
+
+        runner = DistributedRunner(
+            pod_name_prefix="ping",
+            namespace=self.namespace,
+            pod_template=template,
+            timeout_seconds=self.timeout_seconds,
+            log_parse_function=self.parse_pod_logs_from_ping_test_func(),
+        )
+
+        return runner.run_on_all_nodes()

app/test_check_ping.py

@@ -0,0 +1,224 @@
+import unittest
+from unittest.mock import MagicMock, patch
+
+from check_ping import CheckPing
+from pydantic import ValidationError
+
+
+class TestCheckPing(unittest.TestCase):
+    def setUp(self):
+        # Patch the DistributedRunner to avoid actual k8s calls
+        self.mock_distributed_runner_patch = patch("check_ping.DistributedRunner")
+        self.MockDistributedRunner = self.mock_distributed_runner_patch.start()
+        self.mock_runner_instance = MagicMock()
+        self.MockDistributedRunner.return_value = self.mock_runner_instance
+
+        self.valid_params = {
+            "pingTarget": "8.8.8.8",
+        }
+
+    def tearDown(self):
+        self.mock_distributed_runner_patch.stop()
+
+    def test_init_with_minimal_parameters(self):
+        """Tests that CheckPing initializes correctly with minimal valid parameters."""
+        check = CheckPing(self.valid_params)
+        self.assertEqual(check.namespace, "cluster-health")
+        self.assertEqual(check.image, "alpine:3.22")
+        self.assertEqual(check.count, 1)
+        self.assertEqual(check.timeout_seconds, 300)
+        self.assertEqual(check.ping_target, "8.8.8.8")
+        self.assertIsNone(check.secondary_network_config)
+        self.assertIsNone(check.avg_rtt_ms_threshold)
+        self.assertIsNone(check.max_rtt_ms_threshold)
+
+    def test_init_with_custom_parameters(self):
+        """Tests initialization with custom parameters."""
+        params = {
+            "namespace": "custom-ns",
+            "image": "my-custom/ping:latest",
+            "secondaryNetworkConfig": "my-net-config",
+            "timeoutSeconds": 120,
+            "count": 5,
+            "avgRttMsThreshold": 50.0,
+            "maxRttMsThreshold": 100.0,
+            "pingTarget": "1.1.1.1",
+        }
+        check = CheckPing(params)
+        self.assertEqual(check.namespace, "custom-ns")
+        self.assertEqual(check.image, "my-custom/ping:latest")
+        self.assertEqual(check.secondary_network_config, "my-net-config")
+        self.assertEqual(check.timeout_seconds, 120)
+        self.assertEqual(check.count, 5)
+        self.assertEqual(check.avg_rtt_ms_threshold, 50.0)
+        self.assertEqual(check.max_rtt_ms_threshold, 100.0)
+        self.assertEqual(check.ping_target, "1.1.1.1")
+
+    def test_init_with_missing_required_parameters(self):
+        """Tests that initialization fails if required parameters are missing."""
+        with self.assertRaises(ValidationError):
+            CheckPing({"namespace": "test"})  # Missing pingTarget
+
+    # --- Log Parsing Tests ---
+
+    def test_parse_pod_logs_success_no_rtt_check(self):
+        """Tests the log parsing function with successful ping results and no RTT checks."""
+        check = CheckPing(self.valid_params)
+        log_parser = check.parse_pod_logs_from_ping_test_func()
+        logs = """
+        PING 8.8.8.8 (8.8.8.8): 56 data bytes
+        64 bytes from 8.8.8.8: seq=0 ttl=118 time=10.5 ms
+
+        --- 8.8.8.8 ping statistics ---
+        1 packets transmitted, 1 packets received, 0% packet loss
+        """
+        self.assertTrue(log_parser(logs, "node-1"))
+
+    def test_parse_pod_logs_success_with_rtt_check_iputils(self):
+        """Tests log parsing with successful RTTs (iputils-ping format)."""
+        params = self.valid_params.copy()
+        params["avgRttMsThreshold"] = 20.0
+        params["maxRttMsThreshold"] = 30.0
+        check = CheckPing(params)
+        log_parser = check.parse_pod_logs_from_ping_test_func()
+        logs = """
+        PING 8.8.8.8 (8.8.8.8) 56(84) bytes of data.
+        64 bytes from 8.8.8.8: icmp_seq=1 ttl=118 time=15.1 ms
+
+        --- 8.8.8.8 ping statistics ---
+        1 packets transmitted, 1 received, 0% packet loss, time 0ms
+        rtt min/avg/max/mdev = 15.100/15.100/15.100/0.000 ms
+        """
+        self.assertTrue(log_parser(logs, "node-1"))
+
+    def test_parse_pod_logs_success_with_rtt_check_busybox(self):
+        """Tests log parsing with successful RTTs (busybox ping format)."""
+        params = self.valid_params.copy()
+        params["avgRttMsThreshold"] = 20.0
+        params["maxRttMsThreshold"] = 30.0
+        check = CheckPing(params)
+        log_parser = check.parse_pod_logs_from_ping_test_func()
+        logs = """
+        PING 8.8.8.8 (8.8.8.8): 56 data bytes
+        64 bytes from 8.8.8.8: seq=0 ttl=118 time=15.1 ms
+
+        --- 8.8.8.8 ping statistics ---
+        1 packets transmitted, 1 packets received, 0% packet loss
+        round-trip min/avg/max = 15.100/15.100/15.100 ms
+        """
+        self.assertTrue(log_parser(logs, "node-1"))
+
+    def test_parse_pod_logs_packet_loss_fail(self):
+        """Tests log parsing when packet loss is detected."""
+        check = CheckPing(self.valid_params)
+        log_parser = check.parse_pod_logs_from_ping_test_func()
+        logs = "2 packets transmitted, 1 packets received, 50% packet loss"
+        self.assertFalse(log_parser(logs, "node-1"))
+
+    def test_parse_pod_logs_packet_loss_parse_error(self):
+        """Tests log parsing with malformed packet loss data."""
+        check = CheckPing(self.valid_params)
+        log_parser = check.parse_pod_logs_from_ping_test_func()
+        logs = "Some unexpected output without packet loss information."
+        self.assertFalse(log_parser(logs, "node-1"))
+
+    def test_parse_pod_logs_wrong_packet_count(self):
+        """Tests log parsing when the transmitted packet count is wrong."""
+        params = self.valid_params.copy()
+        params["count"] = 5
+        check = CheckPing(params)
+        log_parser = check.parse_pod_logs_from_ping_test_func()
+        logs = "1 packets transmitted, 1 received, 0% packet loss"
+        self.assertFalse(log_parser(logs, "node-1"))
+
+    def test_parse_pod_logs_avg_rtt_fail(self):
+        """Tests log parsing when average RTT is above the threshold."""
+        params = self.valid_params.copy()
+        params["avgRttMsThreshold"] = 10.0
+        check = CheckPing(params)
+        log_parser = check.parse_pod_logs_from_ping_test_func()
+        logs = """
+        1 packets transmitted, 1 received, 0% packet loss
+        rtt min/avg/max/mdev = 15.0/15.0/15.0/0.0 ms
+        """
+        self.assertFalse(log_parser(logs, "node-1"))
+
+    def test_parse_pod_logs_max_rtt_fail(self):
+        """Tests log parsing when max RTT is above the threshold."""
+        params = self.valid_params.copy()
+        params["maxRttMsThreshold"] = 20.0
+        check = CheckPing(params)
+        log_parser = check.parse_pod_logs_from_ping_test_func()
+        logs = """
+        1 packets transmitted, 1 received, 0% packet loss
+        rtt min/avg/max/mdev = 15.0/18.0/25.0/2.0 ms
+        """
+        self.assertFalse(log_parser(logs, "node-1"))
+
+    def test_parse_pod_logs_rtt_parse_error(self):
+        """Tests log parsing with malformed RTT data when RTT check is enabled."""
+        params = self.valid_params.copy()
+        params["avgRttMsThreshold"] = 10.0
+        check = CheckPing(params)
+        log_parser = check.parse_pod_logs_from_ping_test_func()
+        logs = "1 packets transmitted, 1 received, 0% packet loss. No RTT data."
+        self.assertFalse(log_parser(logs, "node-1"))
+
+    # --- is_healthy Tests ---
+
+    def test_is_healthy_success(self):
+        """Tests is_healthy when the distributed run is successful."""
+        self.mock_runner_instance.run_on_all_nodes.return_value = True
+        check = CheckPing(self.valid_params)
+        self.assertTrue(check.is_healthy())
+        self.mock_runner_instance.run_on_all_nodes.assert_called_once()
+
+    def test_is_healthy_failure(self):
+        """Tests is_healthy when the distributed run fails."""
+        self.mock_runner_instance.run_on_all_nodes.return_value = False
+        check = CheckPing(self.valid_params)
+        self.assertFalse(check.is_healthy())
+        self.mock_runner_instance.run_on_all_nodes.assert_called_once()
+
+    def test_is_healthy_constructs_runner_correctly(self):
+        """Tests that DistributedRunner is constructed with the correct pod template."""
+        params = {
+            "pingTarget": "8.8.8.8",
+            "count": 5,
+            "namespace": "ping-ns",
+            "timeoutSeconds": 99,
+        }
+        check = CheckPing(params)
+        check.is_healthy()
+
+        self.MockDistributedRunner.assert_called_once()
+        _, kwargs = self.MockDistributedRunner.call_args
+
+        self.assertEqual(kwargs["pod_name_prefix"], "ping")
+        self.assertEqual(kwargs["namespace"], "ping-ns")
+        self.assertEqual(kwargs["timeout_seconds"], 99)
+        self.assertIsNotNone(kwargs["log_parse_function"])
+
+        pod_template = kwargs["pod_template"]
+        container = pod_template.spec.containers[0]
+        self.assertEqual(container.image, "alpine:3.22")
+        self.assertEqual(container.command, ["ping", "-c", "5", "8.8.8.8"])
+        self.assertIsNone(pod_template.metadata.annotations)
+
+    def test_is_healthy_with_secondary_network(self):
+        """Tests that the pod template includes secondary network annotations."""
+        params = self.valid_params.copy()
+        params["secondaryNetworkConfig"] = '{"name": "net1"}'
+        check = CheckPing(params)
+        check.is_healthy()
+
+        _, kwargs = self.MockDistributedRunner.call_args
+        pod_template = kwargs["pod_template"]
+        self.assertEqual(
+            pod_template.metadata.annotations,
+            {"networking.gke.io/interfaces": '{"name": "net1"}'},
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()