added netprobe arg

Author: jaspals3123

nodescraper/cli/dynamicparserbuilder.py

@@ -24,7 +24,7 @@
 #
 ###############################################################################
 import argparse
-from typing import Optional, Type
+from typing import Literal, Optional, Type, get_args, get_origin
 
 from pydantic import BaseModel
 
@@ -96,19 +96,50 @@ def get_model_arg(cls, type_class_map: dict) -> Optional[Type[BaseModel]]:
             None,
         )
 
+    @classmethod
+    def get_literal_choices(cls, type_class_map: dict) -> Optional[list]:
+        """Get the choices from a Literal type if present
+
+        Args:
+            type_class_map (dict): mapping of type classes
+
+        Returns:
+            Optional[list]: list of valid choices for the Literal type, or None if not a Literal
+        """
+        # Check if Literal is in the type_class_map
+        literal_type = type_class_map.get(Literal)
+        if literal_type and literal_type.inner_type is not None:
+            # The inner_type contains the first literal value, but we need all of them
+            # We need to get the original annotation to extract all literal values
+            # For now, return None and we'll handle this differently
+            return None
+        return None
+
     def add_argument(
         self,
         type_class_map: dict,
         arg_name: str,
         required: bool,
+        annotation: Optional[Type] = None,
     ) -> None:
         """Add an argument to a parser with an appropriate type
 
         Args:
             type_class_map (dict): type classes for the arg
             arg_name (str): argument name
             required (bool): whether or not the arg is required
+            annotation (Optional[Type]): full type annotation for extracting Literal choices
         """
+        # Check for Literal types and extract choices
+        literal_choices = None
+        if Literal in type_class_map and annotation:
+            # Extract all arguments from the annotation
+            args = get_args(annotation)
+            for arg in args:
+                if get_origin(arg) is Literal:
+                    literal_choices = list(get_args(arg))
+                    break
+
         if list in type_class_map:
             type_class = type_class_map[list]
             self.parser.add_argument(
@@ -125,6 +156,15 @@ def add_argument(
                 required=required,
                 choices=[True, False],
             )
+        elif Literal in type_class_map and literal_choices:
+            # Add argument with choices for Literal types
+            self.parser.add_argument(
+                f"--{arg_name}",
+                type=str,
+                required=required,
+                choices=literal_choices,
+                metavar=f"{{{','.join(literal_choices)}}}",
+            )
         elif float in type_class_map:
             self.parser.add_argument(
                 f"--{arg_name}", type=float, required=required, metavar=META_VAR_MAP[float]
@@ -166,6 +206,10 @@ def build_model_arg_parser(self, model: type[BaseModel], required: bool) -> list
             if type(None) in type_class_map and len(attr_data.type_classes) == 1:
                 continue
 
-            self.add_argument(type_class_map, attr.replace("_", "-"), required)
+            # Get the full annotation from the model field
+            field = model.model_fields.get(attr)
+            annotation = field.annotation if field else None
+
+            self.add_argument(type_class_map, attr.replace("_", "-"), required, annotation)
 
         return list(type_map.keys())

nodescraper/plugins/inband/network/collector_args.py

@@ -24,10 +24,11 @@
 #
 ###############################################################################
 
-from typing import Optional
+from typing import Literal, Optional
 
 from nodescraper.models import CollectorArgs
 
 
 class NetworkCollectorArgs(CollectorArgs):
     url: Optional[str] = None
+    netprobe: Optional[Literal["ping", "wget", "curl"]] = None

nodescraper/plugins/inband/network/network_collector.py

@@ -66,6 +66,8 @@ class NetworkCollector(InBandDataCollector[NetworkDataModel, NetworkCollectorArg
     CMD_NEIGHBOR = "ip neighbor show"
     CMD_ETHTOOL_TEMPLATE = "ethtool {interface}"
     CMD_PING = "ping"
+    CMD_WGET = "wget"
+    CMD_CURL = "curl"
 
     # LLDP commands
     CMD_LLDPCLI_NEIGHBOR = "lldpcli show neighbor"
@@ -1671,21 +1673,32 @@ def _collect_pensando_nic_info(
             uncollected_commands,
         )
 
-    def _check_network_connectivity(self, url: str) -> bool:
-        """Check network connectivity by pinging a URL.
+    def _check_network_connectivity(self, cmd: str, url: str) -> bool:
+        """Check network connectivity using specified command.
 
         Args:
-            url: URL or hostname to ping
+            cmd: Command to use for connectivity check (ping, wget, or curl)
+            url: URL or hostname to check
 
         Returns:
             bool: True if network is accessible, False otherwise
         """
+        if cmd not in {"ping", "wget", "curl"}:
+            raise ValueError(
+                f"Invalid network probe command: '{cmd}'. "
+                f"Valid options are: 'ping', 'wget', 'curl'"
+            )
 
         # Determine ping options based on OS
         ping_option = "-c 1" if self.system_info.os_family == OSFamily.LINUX else "-n 1"
 
-        # Run ping command
-        result = self._run_sut_cmd(f"{self.CMD_PING} {url} {ping_option}")
+        # Build command based on cmd parameter using class constants
+        if cmd == "ping":
+            result = self._run_sut_cmd(f"{self.CMD_PING} {url} {ping_option}")
+        elif cmd == "wget":
+            result = self._run_sut_cmd(f"{self.CMD_WGET} {url}")
+        else:  # curl
+            result = self._run_sut_cmd(f"{self.CMD_CURL} {url}")
 
         if result.exit_code == 0:
             self._log_event(
@@ -1697,7 +1710,7 @@ def _check_network_connectivity(self, url: str) -> bool:
         else:
             self._log_event(
                 category=EventCategory.NETWORK,
-                description=f"{self.CMD_PING} to {url} failed!",
+                description=f"{cmd} to {url} failed!",
                 data={"url": url, "not accessible": result.exit_code == 0},
                 priority=EventPriority.ERROR,
             )
@@ -1737,7 +1750,19 @@ def collect_data(
 
         # Check network connectivity if URL is provided
         if args and args.url:
-            network_accessible = self._check_network_connectivity(args.url)
+            cmd = args.netprobe if args.netprobe else "ping"
+            try:
+                network_accessible = self._check_network_connectivity(cmd, args.url)
+            except ValueError as e:
+                self._log_event(
+                    category=EventCategory.NETWORK,
+                    description=str(e),
+                    data={"netprobe": cmd, "url": args.url},
+                    priority=EventPriority.ERROR,
+                    console_log=True,
+                )
+                # Set network_accessible to None since we couldn't check
+                network_accessible = None
 
         # Collect interface/address information
         res_addr = self._run_sut_cmd(self.CMD_ADDR)

test/functional/fixtures/network_plugin_config.json

@@ -3,7 +3,8 @@
   "plugins": {
     "NetworkPlugin": {
       "collector_args": {
-        "url": "mock.example.com"
+        "url": "mock.example.com",
+        "netprobe": "ping"
       },
       "analysis_args": {}
     }

test/functional/test_network_plugin.py

@@ -122,3 +122,135 @@ def test_network_plugin_with_url(run_cli_command, network_config_file, tmp_path)
     assert result.returncode in [0, 1, 2]
     output = result.stdout + result.stderr
     assert len(output) > 0
+
+
+def test_network_plugin_with_netprobe_ping(run_cli_command, tmp_path):
+    """Test NetworkPlugin with netprobe set to ping."""
+    log_path = str(tmp_path / "logs_network_netprobe_ping")
+    result = run_cli_command(
+        [
+            "--log-path",
+            log_path,
+            "run-plugins",
+            "NetworkPlugin",
+            "--url",
+            "google.com",
+            "--netprobe",
+            "ping",
+        ],
+        check=False,
+    )
+
+    assert result.returncode in [0, 1, 2]
+    output = result.stdout + result.stderr
+    assert len(output) > 0
+
+
+def test_network_plugin_with_netprobe_wget(run_cli_command, tmp_path):
+    """Test NetworkPlugin with netprobe set to wget."""
+    log_path = str(tmp_path / "logs_network_netprobe_wget")
+    result = run_cli_command(
+        [
+            "--log-path",
+            log_path,
+            "run-plugins",
+            "NetworkPlugin",
+            "--url",
+            "google.com",
+            "--netprobe",
+            "wget",
+        ],
+        check=False,
+    )
+
+    assert result.returncode in [0, 1, 2]
+    output = result.stdout + result.stderr
+    assert len(output) > 0
+
+
+def test_network_plugin_with_netprobe_curl(run_cli_command, tmp_path):
+    """Test NetworkPlugin with netprobe set to curl."""
+    log_path = str(tmp_path / "logs_network_netprobe_curl")
+    result = run_cli_command(
+        [
+            "--log-path",
+            log_path,
+            "run-plugins",
+            "NetworkPlugin",
+            "--url",
+            "google.com",
+            "--netprobe",
+            "curl",
+        ],
+        check=False,
+    )
+
+    assert result.returncode in [0, 1, 2]
+    output = result.stdout + result.stderr
+    assert len(output) > 0
+
+
+def test_network_plugin_with_invalid_netprobe(run_cli_command, tmp_path):
+    """Test NetworkPlugin with invalid netprobe value - should fail at CLI validation."""
+    log_path = str(tmp_path / "logs_network_invalid_netprobe")
+    result = run_cli_command(
+        [
+            "--log-path",
+            log_path,
+            "run-plugins",
+            "NetworkPlugin",
+            "--url",
+            "google.com",
+            "--netprobe",
+            "invalid",
+        ],
+        check=False,
+    )
+
+    # Should fail with exit code 2 (argparse error)
+    assert result.returncode == 2
+    output = result.stdout + result.stderr
+    assert len(output) > 0
+    assert "invalid choice" in output.lower()
+    assert "choose from" in output.lower()
+
+
+def test_network_plugin_with_url_no_netprobe(run_cli_command, tmp_path):
+    """Test NetworkPlugin with URL but no netprobe - should default to ping."""
+    log_path = str(tmp_path / "logs_network_url_default")
+    result = run_cli_command(
+        [
+            "--log-path",
+            log_path,
+            "run-plugins",
+            "NetworkPlugin",
+            "--url",
+            "google.com",
+        ],
+        check=False,
+    )
+
+    assert result.returncode in [0, 1, 2]
+    output = result.stdout + result.stderr
+    assert len(output) > 0
+
+
+def test_network_plugin_with_netprobe_no_url(run_cli_command, tmp_path):
+    """Test NetworkPlugin with netprobe but no URL - should skip connectivity check."""
+    log_path = str(tmp_path / "logs_network_netprobe_no_url")
+    result = run_cli_command(
+        [
+            "--log-path",
+            log_path,
+            "run-plugins",
+            "NetworkPlugin",
+            "--netprobe",
+            "ping",
+        ],
+        check=False,
+    )
+
+    # Should succeed but skip connectivity check
+    assert result.returncode in [0, 1, 2]
+    output = result.stdout + result.stderr
+    assert len(output) > 0