ANTA Tests for VXLAN

class VerifyVxlan1ConnSettings(AntaTest):
    """Verifies Vxlan1 source interface and UDP port.

    Expected Results
    ----------------
    * Success: Passes if the interface vxlan1 source interface and UDP port are correct.
    * Failure: Fails if the interface vxlan1 source interface or UDP port are incorrect.
    * Skipped: Skips if the Vxlan1 interface is not configured.

    Examples
    --------
    ```yaml
    anta.tests.vxlan:
      - VerifyVxlan1ConnSettings:
          source_interface: Loopback1
          udp_port: 4789
    ```
    """

    categories: ClassVar[list[str]] = ["vxlan"]
    commands: ClassVar[list[AntaCommand | AntaTemplate]] = [AntaCommand(command="show interfaces", revision=1)]

    class Input(AntaTest.Input):
        """Input model for the VerifyVxlan1ConnSettings test."""

        source_interface: VxlanSrcIntf
        """Source interface of vxlan1 interface."""
        udp_port: int = Field(ge=1024, le=65335)
        """UDP port used for vxlan1 interface."""

    @AntaTest.anta_test
    def test(self) -> None:
        """Main test function for VerifyVxlan1ConnSettings."""
        self.result.is_success()
        command_output = self.instance_commands[0].json_output

        # Skip the test case if vxlan1 interface is not configured
        vxlan_output = get_value(command_output, "interfaces.Vxlan1")
        if not vxlan_output:
            self.result.is_skipped("Interface: Vxlan1 - Not configured")
            return

        src_intf = vxlan_output.get("srcIpIntf")
        port = vxlan_output.get("udpPort")

        # Check vxlan1 source interface and udp port
        if src_intf != self.inputs.source_interface:
            self.result.is_failure(f"Interface: Vxlan1 - Incorrect Source interface - Expected: {self.inputs.source_interface} Actual: {src_intf}")
        if port != self.inputs.udp_port:
            self.result.is_failure(f"Interface: Vxlan1 - Incorrect UDP port - Expected: {self.inputs.udp_port} Actual: {port}")

class VerifyVxlan1Interface(AntaTest):
    """Verifies the Vxlan1 interface status.

    Warnings
    --------
    The name of this test has been updated from 'VerifyVxlan' for better representation.

    Expected Results
    ----------------
    * Success: The test will pass if the Vxlan1 interface is configured with line protocol status and interface status 'up'.
    * Failure: The test will fail if the Vxlan1 interface line protocol status or interface status are not 'up'.
    * Skipped: The test will be skipped if the Vxlan1 interface is not configured.

    Examples
    --------
    ```yaml
    anta.tests.vxlan:
      - VerifyVxlan1Interface:
    ```
    """

    categories: ClassVar[list[str]] = ["vxlan"]
    commands: ClassVar[list[AntaCommand | AntaTemplate]] = [AntaCommand(command="show interfaces description", revision=1)]

    @AntaTest.anta_test
    def test(self) -> None:
        """Main test function for VerifyVxlan1Interface."""
        self.result.is_success()
        command_output = self.instance_commands[0].json_output

        # Skipping the test if the Vxlan1 interface is not configured
        if "Vxlan1" not in (interface_details := command_output["interfaceDescriptions"]):
            self.result.is_skipped("Interface: Vxlan1 - Not configured")
            return

        line_protocol_status = interface_details["Vxlan1"]["lineProtocolStatus"]
        interface_status = interface_details["Vxlan1"]["interfaceStatus"]

        # Checking against both status and line protocol status
        if interface_status != "up" or line_protocol_status != "up":
            self.result.is_failure(f"Interface: Vxlan1 - Incorrect Line protocol status/Status - Expected: up/up Actual: {line_protocol_status}/{interface_status}")

class VerifyVxlan1VVTEPIPAddresses(AntaTest):
    """Verifies the VVTEP IP addresses.

    Supports IPv4-only, IPv6-only, or dual-stack. At least one address must be provided.

    Expected Results
    ----------------
    * Success: Passes if the provided VVTEP IP addresses match the ones configured on the VXLAN interface.
    * Failure: Fails if any provided VVTEP IP address does not match the configured value, or is not configured.
    * Skipped: Skips if the Vxlan1 interface is not configured.

    Examples
    --------
    ```yaml
    anta.tests.vxlan:
      - VerifyVxlan1VVTEPIPAddresses:
          ipv4_address: 5.5.5.5
          ipv6_address: fd00:dc:1::1
    ```
    """

    categories: ClassVar[list[str]] = ["vxlan"]
    commands: ClassVar[list[AntaCommand | AntaTemplate]] = [AntaCommand(command="show interfaces", revision=1)]

    class Input(AntaTest.Input):
        """Input model for the VerifyVxlan1VVTEPIPAddresses test."""

        ipv4_address: IPv4Address | None = None
        """VVTEP IPv4 address."""
        ipv6_address: IPv6Address | None = None
        """VVTEP IPv6 address."""

        @model_validator(mode="after")
        def validate_inputs(self) -> Self:
            """Validate that at least one address is provided."""
            if not self.ipv4_address and not self.ipv6_address:
                msg = "'ipv4_address' or 'ipv6_address' must be provided"
                raise ValueError(msg)
            return self

    @AntaTest.anta_test
    def test(self) -> None:
        """Main test function for VerifyVxlan1VVTEPIPAddresses."""
        self.result.is_success()
        command_output = self.instance_commands[0].json_output

        # Skip the test case if vxlan1 interface is not configured
        vxlan_output = get_value(command_output, "interfaces.Vxlan1")
        if not vxlan_output:
            self.result.is_skipped("Interface: Vxlan1 - Not configured")
            return

        checks: list[tuple[str, str, str]] = []
        if self.inputs.ipv4_address is not None:
            checks.append((str(self.inputs.ipv4_address), "vArpVtepAddr", "IPv4"))
        if self.inputs.ipv6_address is not None:
            checks.append((str(self.inputs.ipv6_address), "vArpVtepAddrV6", "IPv6"))

        for expected, key, ip_version in checks:
            configured = vxlan_output.get(key)
            if configured is None:
                self.result.is_failure(f"Interface: Vxlan1 - VVTEP {ip_version} address is not configured")
            elif expected != configured:
                self.result.is_failure(f"Interface: Vxlan1 - Incorrect VVTEP {ip_version} address - Expected: {expected} Actual: {configured}")

class VerifyVxlanConfigSanity(AntaTest):
    """Verifies there are no VXLAN config-sanity inconsistencies.

    Expected Results
    ----------------
    * Success: The test will pass if no issues are detected with the VXLAN configuration.
    * Failure: The test will fail if issues are detected with the VXLAN configuration.
    * Skipped: The test will be skipped if VXLAN is not configured on the device.

    Examples
    --------
    ```yaml
    anta.tests.vxlan:
      - VerifyVxlanConfigSanity:
    ```
    """

    categories: ClassVar[list[str]] = ["vxlan"]
    commands: ClassVar[list[AntaCommand | AntaTemplate]] = [AntaCommand(command="show vxlan config-sanity", revision=1)]

    @AntaTest.anta_test
    def test(self) -> None:
        """Main test function for VerifyVxlanConfigSanity."""
        self.result.is_success()
        command_output = self.instance_commands[0].json_output

        # Skipping the test if VXLAN is not configured
        if "categories" not in command_output or len(command_output["categories"]) == 0:
            self.result.is_skipped("VXLAN is not configured")
            return

        # Verifies the Vxlan config sanity
        categories_to_check = ["localVtep", "mlag", "pd"]
        for category in categories_to_check:
            if not get_value(command_output, f"categories.{category}.allCheckPass"):
                self.result.is_failure(f"Vxlan Category: {category} - Config sanity check is not passing")

class VerifyVxlanVniBinding(AntaTest):
    """Verifies the VNI-VLAN, VNI-VRF bindings of the Vxlan1 interface.

    Expected Results
    ----------------
    * Success: The test will pass if the VNI-VLAN and VNI-VRF bindings provided are properly configured.
    * Failure: The test will fail if any VNI lacks bindings or if any bindings are incorrect.
    * Skipped: The test will be skipped if the Vxlan1 interface is not configured.

    Examples
    --------
    ```yaml
    anta.tests.vxlan:
      - VerifyVxlanVniBinding:
          bindings:
            10010: 10
            10020: 20
            500: PROD
    ```
    """

    categories: ClassVar[list[str]] = ["vxlan"]
    commands: ClassVar[list[AntaCommand | AntaTemplate]] = [AntaCommand(command="show vxlan vni", revision=1)]

    class Input(AntaTest.Input):
        """Input model for the VerifyVxlanVniBinding test."""

        bindings: dict[Vni, VlanId | str]
        """VNI-VLAN or VNI-VRF bindings to verify."""

    @AntaTest.anta_test
    def test(self) -> None:
        """Main test function for VerifyVxlanVniBinding."""
        self.result.is_success()

        if (vxlan1 := get_value(self.instance_commands[0].json_output, "vxlanIntfs.Vxlan1")) is None:
            self.result.is_skipped("Interface: Vxlan1 - Not configured")
            return

        for vni, vlan_vrf in self.inputs.bindings.items():
            str_vni = str(vni)
            retrieved_vlan = ""
            retrieved_vrf = ""
            if all([str_vni in vxlan1["vniBindings"], isinstance(vlan_vrf, int)]):
                retrieved_vlan = get_value(vxlan1, f"vniBindings..{str_vni}..vlan", separator="..")
            elif str_vni in vxlan1["vniBindingsToVrf"]:
                if isinstance(vlan_vrf, int):
                    retrieved_vlan = get_value(vxlan1, f"vniBindingsToVrf..{str_vni}..vlan", separator="..")
                else:
                    retrieved_vrf = get_value(vxlan1, f"vniBindingsToVrf..{str_vni}..vrfName", separator="..")
            if not any([retrieved_vlan, retrieved_vrf]):
                self.result.is_failure(f"Interface: Vxlan1 VNI: {str_vni} - Binding not found")

            elif retrieved_vlan and vlan_vrf != retrieved_vlan:
                self.result.is_failure(f"Interface: Vxlan1 VNI: {str_vni} - Wrong VLAN binding - Expected: {vlan_vrf} Actual: {retrieved_vlan}")

            elif retrieved_vrf and vlan_vrf != retrieved_vrf:
                self.result.is_failure(f"Interface: Vxlan1 VNI: {str_vni} - Wrong VRF binding - Expected: {vlan_vrf} Actual: {retrieved_vrf}")

class VerifyVxlanVtep(AntaTest):
    """Verifies Vxlan1 VTEP peers.

    Expected Results
    ----------------
    * Success: The test will pass if all provided VTEP peers are identified and matching.
    * Failure: The test will fail if any VTEP peer is missing or there are unexpected VTEP peers.
    * Skipped: The test will be skipped if the Vxlan1 interface is not configured.

    Examples
    --------
    ```yaml
    anta.tests.vxlan:
      - VerifyVxlanVtep:
          vteps:
            - 10.1.1.5
            - 10.1.1.6
    ```
    """

    categories: ClassVar[list[str]] = ["vxlan"]
    commands: ClassVar[list[AntaCommand | AntaTemplate]] = [AntaCommand(command="show vxlan vtep", revision=1)]

    class Input(AntaTest.Input):
        """Input model for the VerifyVxlanVtep test."""

        vteps: list[IPv4Address]
        """List of VTEP peers to verify."""

    @AntaTest.anta_test
    def test(self) -> None:
        """Main test function for VerifyVxlanVtep."""
        self.result.is_success()

        inputs_vteps = [str(input_vtep) for input_vtep in self.inputs.vteps]

        if (vxlan1 := get_value(self.instance_commands[0].json_output, "interfaces.Vxlan1")) is None:
            self.result.is_skipped("Interface: Vxlan1 - Not configured")
            return

        difference1 = set(inputs_vteps).difference(set(vxlan1["vteps"]))
        difference2 = set(vxlan1["vteps"]).difference(set(inputs_vteps))

        if difference1:
            self.result.is_failure(f"The following VTEP peer(s) are missing from the Vxlan1 interface: {', '.join(sorted(difference1))}")

        if difference2:
            self.result.is_failure(f"Unexpected VTEP peer(s) on Vxlan1 interface: {', '.join(sorted(difference2))}")