Repository: aristanetworks/anta (python)

Issue:
Support JSON for inventory
Similar to #714, ANTA should be able to load an inventory file using `json` format

Pull request:
feat(anta): Support JSON inventory
# Description

cf title

Fixes #742 

# Checklist:

- [x] My code follows the style guidelines of this project
- [x] I have performed a self-review of my code
- [x] I have run pre-commit for code linting and typing (`pre-commit run`)
- [x] I have made corresponding changes to the documentation
- [x] I have added tests that prove my fix is effective or that my feature works
- [x] New and existing unit tests pass locally with my changes (`tox -e testenv`)


Patch:
diff --git a/anta/cli/utils.py b/anta/cli/utils.py
index 508424dd..34b96b3f 100644
--- a/anta/cli/utils.py
+++ b/anta/cli/utils.py
@@ -9,7 +9,7 @@ import enum
 import functools
 import logging
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Callable
+from typing import TYPE_CHECKING, Any, Callable, Literal
 
 import click
 from yaml import YAMLError
@@ -191,6 +191,14 @@ def core_options(f: Callable[..., Any]) -> Callable[..., Any]:
         required=True,
         type=click.Path(file_okay=True, dir_okay=False, exists=True, readable=True, path_type=Path),
     )
+    @click.option(
+        "--inventory-format",
+        envvar="ANTA_INVENTORY_FORMAT",
+        show_envvar=True,
+        help="Format of the inventory file, either 'yaml' or 'json'",
+        default="yaml",
+        type=click.Choice(["yaml", "json"], case_sensitive=False),
+    )
     @click.pass_context
     @functools.wraps(f)
     def wrapper(
@@ -205,6 +213,7 @@ def core_options(f: Callable[..., Any]) -> Callable[..., Any]:
         timeout: float,
         insecure: bool,
         disable_cache: bool,
+        inventory_format: Literal["json", "yaml"],
         **kwargs: dict[str, Any],
     ) -> Any:
         # If help is invoke somewhere, do not parse inventory
@@ -242,6 +251,7 @@ def core_options(f: Callable[..., Any]) -> Callable[..., Any]:
                 timeout=timeout,
                 insecure=insecure,
                 disable_cache=disable_cache,
+                file_format=inventory_format,
             )
         except (TypeError, ValueError, YAMLError, OSError, InventoryIncorrectSchemaError, InventoryRootKeyError) as e:
             anta_log_exception(e, f"Failed to parse the inventory: {inventory}", logger)
diff --git a/anta/inventory/__init__.py b/anta/inventory/__init__.py
index f98c42f2..a62afacc 100644
--- a/anta/inventory/__init__.py
+++ b/anta/inventory/__init__.py
@@ -8,8 +8,9 @@ from __future__ import annotations
 import asyncio
 import logging
 from ipaddress import ip_address, ip_network
+from json import load as json_load
 from pathlib import Path
-from typing import Any, ClassVar
+from typing import Any, ClassVar, Literal
 
 from pydantic import ValidationError
 from yaml import YAMLError, safe_load
@@ -178,6 +179,7 @@ class AntaInventory(dict[str, AntaDevice]):
         password: str,
         enable_password: str | None = None,
         timeout: float | None = None,
+        file_format: Literal["yaml", "json"] = "yaml",
         *,
         enable: bool = False,
         insecure: bool = False,
@@ -199,6 +201,8 @@ class AntaInventory(dict[str, AntaDevice]):
             Enable password to use if required.
         timeout
             Timeout value in seconds for outgoing API calls.
+        file_format
+            Whether the inventory file is in JSON or YAML.
         enable
             Whether or not the commands need to be run in enable mode towards the devices.
         insecure
@@ -214,6 +218,10 @@ class AntaInventory(dict[str, AntaDevice]):
             Inventory file is not following AntaInventory Schema.
 
         """
+        if file_format not in ["yaml", "json"]:
+            message = f"'{file_format}' is not a valid format for an AntaInventory file. Only 'yaml' and 'json' are supported."
+            raise ValueError(message)
+
         inventory = AntaInventory()
         kwargs: dict[str, Any] = {
             "username": username,
@@ -224,20 +232,12 @@ class AntaInventory(dict[str, AntaDevice]):
             "insecure": insecure,
             "disable_cache": disable_cache,
         }
-        if username is None:
-            message = "'username' is required to create an AntaInventory"
-            logger.error(message)
-            raise ValueError(message)
-        if password is None:
-            message = "'password' is required to create an AntaInventory"
-            logger.error(message)
-            raise ValueError(message)
 
         try:
             filename = Path(filename)
             with filename.open(encoding="UTF-8") as file:
-                data = safe_load(file)
-        except (TypeError, YAMLError, OSError) as e:
+                data = safe_load(file) if file_format == "yaml" else json_load(file)
+        except (TypeError, YAMLError, OSError, ValueError) as e:
             message = f"Unable to parse ANTA Device Inventory file '{filename}'"
             anta_log_exception(e, message, logger)
             raise
diff --git a/anta/inventory/models.py b/anta/inventory/models.py
index 320f5276..493bad71 100644
--- a/anta/inventory/models.py
+++ b/anta/inventory/models.py
@@ -9,14 +9,26 @@ import logging
 import math
 
 import yaml
-from pydantic import BaseModel, ConfigDict, IPvAnyAddress, IPvAnyNetwork
+from pydantic import BaseModel, ConfigDict, FieldSerializationInfo, IPvAnyAddress, IPvAnyNetwork, field_serializer
 
 from anta.custom_types import Hostname, Port
 
 logger = logging.getLogger(__name__)
 
 
-class AntaInventoryHost(BaseModel):
+class AntaInventoryBaseModel(BaseModel):
+    """Pydantic BaseModel for AntaInventory objects."""
+
+    model_config = ConfigDict(extra="forbid")
+
+    # Using check_fields as we plan to use this in the child classes
+    @field_serializer("tags", when_used="json", check_fields=False)
+    def serialize_tags(self, tags: set[str], _info: FieldSerializationInfo) -> list[str]:
+        """Make sure the tags are always dumped in the same order."""
+        return sorted(tags)
+
+
+class AntaInventoryHost(AntaInventoryBaseModel):
     """Host entry of AntaInventoryInput.
 
     Attributes
@@ -34,8 +46,6 @@ class AntaInventoryHost(BaseModel):
 
     """
 
-    model_config = ConfigDict(extra="forbid")
-
     name: str | None = None
     host: Hostname | IPvAnyAddress
     port: Port | None = None
@@ -43,7 +53,7 @@ class AntaInventoryHost(BaseModel):
     disable_cache: bool = False
 
 
-class AntaInventoryNetwork(BaseModel):
+class AntaInventoryNetwork(AntaInventoryBaseModel):
     """Network entry of AntaInventoryInput.
 
     Attributes
@@ -57,14 +67,12 @@ class AntaInventoryNetwork(BaseModel):
 
     """
 
-    model_config = ConfigDict(extra="forbid")
-
     network: IPvAnyNetwork
     tags: set[str] | None = None
     disable_cache: bool = False
 
 
-class AntaInventoryRange(BaseModel):
+class AntaInventoryRange(AntaInventoryBaseModel):
     """IP Range entry of AntaInventoryInput.
 
     Attributes
@@ -80,8 +88,6 @@ class AntaInventoryRange(BaseModel):
 
     """
 
-    model_config = ConfigDict(extra="forbid")
-
     start: IPvAnyAddress
     end: IPvAnyAddress
     tags: set[str] | None = None
@@ -109,4 +115,13 @@ class AntaInventoryInput(BaseModel):
         # This could be improved.
         # https://github.com/pydantic/pydantic/issues/1043
         # Explore if this worth using this: https://github.com/NowanIlfideme/pydantic-yaml
-        return yaml.safe_dump(yaml.safe_load(self.model_dump_json(serialize_as_any=True, exclude_unset=True)), indent=2, width=math.inf)
+        return yaml.safe_dump(yaml.safe_load(self.model_dump_json(serialize_as_any=True, exclude_unset=True)), width=math.inf)
+
+    def to_json(self) -> str:
+        """Return a JSON representation string of this model.
+
+        Returns
+        -------
+            The JSON representation string of this model.
+        """
+        return self.model_dump_json(serialize_as_any=True, exclude_unset=True, indent=2)
diff --git a/docs/snippets/anta_nrfu_help.txt b/docs/snippets/anta_nrfu_help.txt
index 5801f4e3..5500389b 100644
--- a/docs/snippets/anta_nrfu_help.txt
+++ b/docs/snippets/anta_nrfu_help.txt
@@ -28,6 +28,8 @@ Options:
                                   ANTA_DISABLE_CACHE]
   -i, --inventory FILE            Path to the inventory YAML file.  [env var:
                                   ANTA_INVENTORY; required]
+  --inventory-format [yaml|json]  Format of the inventory file, either 'yaml'
+                                  or 'json'  [env var: ANTA_INVENTORY_FORMAT]
   --tags TEXT                     List of tags using comma as separator:
                                   tag1,tag2,tag3.  [env var: ANTA_TAGS]
   -c, --catalog FILE              Path to the test catalog file  [env var:


Test patch:
diff --git a/tests/data/test_inventory_large.json b/tests/data/test_inventory_large.json
new file mode 100644
index 00000000..876d9409
--- /dev/null
+++ b/tests/data/test_inventory_large.json
@@ -0,0 +1,207 @@
+{
+  "anta_inventory": {
+    {
+      "hosts": [
+        {
+          "name": "super-spine1",
+          "host": "localhost"
+        },
+        {
+          "name": "super-spine2",
+          "host": "localhost"
+        },
+        {
+          "name": "pod1-spine1",
+          "host": "localhost"
+        },
+        {
+          "name": "pod1-spine2",
+          "host": "localhost"
+        },
+        {
+          "name": "pod1-spine3",
+          "host": "localhost"
+        },
+        {
+          "name": "pod1-spine4",
+          "host": "localhost"
+        },
+        {
+          "name": "pod1-leaf1a",
+          "host": "localhost"
+        },
+        {
+          "name": "pod1-leaf1b",
+          "host": "localhost"
+        },
+        {
+          "name": "pod1-leaf2a",
+          "host": "localhost"
+        },
+        {
+          "name": "pod1-leaf2b",
+          "host": "localhost"
+        },
+        {
+          "name": "pod1-leaf3a",
+          "host": "localhost"
+        },
+        {
+          "name": "pod1-leaf3b",
+          "host": "localhost"
+        },
+        {
+          "name": "pod1-leaf4a",
+          "host": "localhost"
+        },
+        {
+          "name": "pod1-leaf4b",
+          "host": "localhost"
+        },
+        {
+          "name": "pod1-leaf5a",
+          "host": "localhost"
+        },
+        {
+          "name": "pod1-leaf5b",
+          "host": "localhost"
+        },
+        {
+          "name": "pod1-leaf6a",
+          "host": "localhost"
+        },
+        {
+          "name": "pod1-leaf6b",
+          "host": "localhost"
+        },
+        {
+          "name": "pod1-leaf7a",
+          "host": "localhost"
+        },
+        {
+          "name": "pod1-leaf7b",
+          "host": "localhost"
+        },
+        {
+          "name": "pod1-leaf8a",
+          "host": "localhost"
+        },
+        {
+          "name": "pod1-leaf8b",
+          "host": "localhost"
+        },
+        {
+          "name": "pod1-leaf9a",
+          "host": "localhost"
+        },
+        {
+          "name": "pod1-leaf9b",
+          "host": "localhost"
+        },
+        {
+          "name": "pod1-leaf10a",
+          "host": "localhost"
+        },
+        {
+          "name": "pod1-leaf10b",
+          "host": "localhost"
+        },
+        {
+          "name": "pod2-spine1",
+          "host": "localhost"
+        },
+        {
+          "name": "pod2-spine2",
+          "host": "localhost"
+        },
+        {
+          "name": "pod2-spine3",
+          "host": "localhost"
+        },
+        {
+          "name": "pod2-spine4",
+          "host": "localhost"
+        },
+        {
+          "name": "pod2-leaf1a",
+          "host": "localhost"
+        },
+        {
+          "name": "pod2-leaf1b",
+          "host": "localhost"
+        },
+        {
+          "name": "pod2-leaf2a",
+          "host": "localhost"
+        },
+        {
+          "name": "pod2-leaf2b",
+          "host": "localhost"
+        },
+        {
+          "name": "pod2-leaf3a",
+          "host": "localhost"
+        },
+        {
+          "name": "pod2-leaf3b",
+          "host": "localhost"
+        },
+        {
+          "name": "pod2-leaf4a",
+          "host": "localhost"
+        },
+        {
+          "name": "pod2-leaf4b",
+          "host": "localhost"
+        },
+        {
+          "name": "pod2-leaf5a",
+          "host": "localhost"
+        },
+        {
+          "name": "pod2-leaf5b",
+          "host": "localhost"
+        },
+        {
+          "name": "pod2-leaf6a",
+          "host": "localhost"
+        },
+        {
+          "name": "pod2-leaf6b",
+          "host": "localhost"
+        },
+        {
+          "name": "pod2-leaf7a",
+          "host": "localhost"
+        },
+        {
+          "name": "pod2-leaf7b",
+          "host": "localhost"
+        },
+        {
+          "name": "pod2-leaf8a",
+          "host": "localhost"
+        },
+        {
+          "name": "pod2-leaf8b",
+          "host": "localhost"
+        },
+        {
+          "name": "pod2-leaf9a",
+          "host": "localhost"
+        },
+        {
+          "name": "pod2-leaf9b",
+          "host": "localhost"
+        },
+        {
+          "name": "pod2-leaf10a",
+          "host": "localhost"
+        },
+        {
+          "name": "pod2-leaf10b",
+          "host": "localhost"
+        }
+      ]
+    }
+}
diff --git a/tests/data/test_inventory_medium.json b/tests/data/test_inventory_medium.json
new file mode 100644
index 00000000..5c4f1473
--- /dev/null
+++ b/tests/data/test_inventory_medium.json
@@ -0,0 +1,30 @@
+{
+  "anta_inventory": {
+    "hosts": [
+      {
+        "name": "spine1",
+        "host": "localhost"
+      },
+      {
+        "name": "spine2",
+        "host": "localhost"
+      },
+      {
+        "name": "leaf1a",
+        "host": "localhost"
+      },
+      {
+        "name": "leaf1b",
+        "host": "localhost"
+      },
+      {
+        "name": "leaf2a",
+        "host": "localhost"
+      },
+      {
+        "name": "leaf2b",
+        "host": "localhost"
+      }
+    ]
+  }
+}
diff --git a/tests/data/test_inventory_with_tags.json b/tests/data/test_inventory_with_tags.json
new file mode 100644
index 00000000..0f0aa7ce
--- /dev/null
+++ b/tests/data/test_inventory_with_tags.json
@@ -0,0 +1,22 @@
+{
+  "anta_inventory": {
+    "hosts": [
+      {
+        "name": "leaf1",
+        "host": "leaf1.anta.arista.com",
+        "tags": ["dc1", "leaf"]
+      },
+      {
+        "name": "leaf2",
+        "host": "leaf2.anta.arista.com",
+        "tags": ["leaf"]
+      },
+      {
+        "name": "spine1",
+        "host": "spine1.anta.arista.com",
+        "tags": ["spine"],
+        "disable_cache": true
+      }
+    ]
+  }
+}
diff --git a/tests/data/test_inventory_with_tags.yml b/tests/data/test_inventory_with_tags.yml
index 16a9df4c..a3eb0c00 100644
--- a/tests/data/test_inventory_with_tags.yml
+++ b/tests/data/test_inventory_with_tags.yml
@@ -3,10 +3,11 @@ anta_inventory:
   hosts:
     - name: leaf1
       host: leaf1.anta.arista.com
-      tags: ["leaf", "dc1"]
+      tags: ["dc1", "leaf"]
     - name: leaf2
       host: leaf2.anta.arista.com
       tags: ["leaf"]
     - name: spine1
       host: spine1.anta.arista.com
       tags: ["spine"]
+      disable_cache: true
diff --git a/tests/units/inventory/test__init__.py b/tests/units/inventory/test__init__.py
index 9eee8803..461f0379 100644
--- a/tests/units/inventory/test__init__.py
+++ b/tests/units/inventory/test__init__.py
@@ -5,7 +5,7 @@
 
 from __future__ import annotations
 
-from pathlib import Path
+import logging
 from typing import TYPE_CHECKING
 
 import pytest
@@ -15,10 +15,9 @@ from anta.inventory import AntaInventory
 from anta.inventory.exceptions import InventoryIncorrectSchemaError, InventoryRootKeyError
 
 if TYPE_CHECKING:
-    from _pytest.mark.structures import ParameterSet
-
-FILE_DIR: Path = Path(__file__).parent.parent.resolve() / "data" / "inventory"
+    from pathlib import Path
 
+    from _pytest.mark.structures import ParameterSet
 
 INIT_VALID_PARAMS: list[ParameterSet] = [
     pytest.param(
@@ -76,3 +75,15 @@ class TestAntaInventory:
         """Parse invalid YAML file to create ANTA inventory."""
         with pytest.raises((InventoryIncorrectSchemaError, InventoryRootKeyError, ValidationError)):
             AntaInventory.parse(filename=yaml_file, username="arista", password="arista123")
+
+    def test_parse_wrong_format(self) -> None:
+        """Use wrong file format to parse the ANTA inventory."""
+        with pytest.raises(ValueError, match=" is not a valid format for an AntaInventory file. Only 'yaml' and 'json' are supported."):
+            AntaInventory.parse(filename="dummy.yml", username="arista", password="arista123", file_format="wrong")  # type: ignore[arg-type]
+
+    def test_parse_os_error(self, caplog: pytest.LogCaptureFixture) -> None:
+        """Use wrong file name to parse the ANTA inventory."""
+        caplog.set_level(logging.INFO)
+        with pytest.raises(OSError, match="No such file or directory"):
+            _ = AntaInventory.parse(filename="dummy.yml", username="arista", password="arista123")
+        assert "Unable to parse ANTA Device Inventory file" in caplog.records[0].message
diff --git a/tests/units/inventory/test_models.py b/tests/units/inventory/test_models.py
index 63ce7a78..f6bd7d63 100644
--- a/tests/units/inventory/test_models.py
+++ b/tests/units/inventory/test_models.py
@@ -5,16 +5,21 @@
 
 from __future__ import annotations
 
+import json
+from pathlib import Path
 from typing import TYPE_CHECKING, Any
 
 import pytest
 from pydantic import ValidationError
+from yaml import safe_load
 
-from anta.inventory.models import AntaInventoryHost, AntaInventoryNetwork, AntaInventoryRange
+from anta.inventory.models import AntaInventoryHost, AntaInventoryInput, AntaInventoryNetwork, AntaInventoryRange
 
 if TYPE_CHECKING:
     from _pytest.mark.structures import ParameterSet
 
+FILE_DIR: Path = Path(__file__).parents[2] / "data"
+
 INVENTORY_HOST_VALID_PARAMS: list[ParameterSet] = [
     pytest.param(None, "1.1.1.1", None, None, None, id="IPv4"),
     pytest.param(None, "fe80::cc62:a9ff:feef:932a", None, None, None, id="IPv6"),
@@ -164,3 +169,33 @@ class TestAntaInventoryRange:
         """Invalid model parameters."""
         with pytest.raises(ValidationError):
             AntaInventoryRange.model_validate({"start": start, "end": end, "tags": tags, "disable_cache": disable_cache})
+
+
+class TestAntaInventoryInputs:
+    """Test anta.inventory.models.AntaInventoryInputs."""
+
+    def test_dump_to_json(self) -> None:
+        """Load a YAML file, dump it to JSON and verify it works."""
+        input_yml_path = FILE_DIR / "test_inventory_with_tags.yml"
+        expected_json_path = FILE_DIR / "test_inventory_with_tags.json"
+        with input_yml_path.open("r") as f:
+            data = safe_load(f)
+        anta_inventory_input = AntaInventoryInput(**data["anta_inventory"])
+
+        with expected_json_path.open("r") as f:
+            expected_data = json.load(f)
+
+        assert anta_inventory_input.to_json() == json.dumps(expected_data["anta_inventory"], indent=2)
+
+    def test_dump_to_yaml(self) -> None:
+        """Load a JSON file, dump it to YAML and verify it works."""
+        input_json_path = FILE_DIR / "test_inventory_medium.json"
+        expected_yml_path = FILE_DIR / "test_inventory_medium.yml"
+        with input_json_path.open("r") as f:
+            data = json.load(f)
+        anta_inventory_input = AntaInventoryInput(**data["anta_inventory"])
+
+        with expected_yml_path.open("r") as f:
+            expected_data = safe_load(f)
+
+        assert safe_load(anta_inventory_input.yaml()) == expected_data["anta_inventory"]