Merge pull request #106 from network-unit-testing-system/105-poor-performance-on-big-test-portfolio-using-remote-nornir-inventory

add option to not cache the inventory

Author: marcom4rtinez

docs/source/advanced/cli.rst

@@ -0,0 +1,28 @@
+Advanced Options
+================
+
+Nuts can be tweaked if the default options are not sufficient.
+
+Custom Nornir Inventory
+-----------------------
+
+If you already have a nornir configuration and inventory for your network you can reuse it by passing the parameter ``--nornir-config`` to the pytest command:
+
+.. code:: shell
+
+    $ pytest tests/test-definition-ping.yaml --nornir-config path/to/nr-config.yaml
+
+
+It will default back to the local directory.
+
+
+Disable caching of the Nornir Inventory
+---------------------------------------
+
+For performance reasons the Nornir Inventory will be internally cached. If you wish to disable this behavior pass the parameter ``--nornir-cache-disable`` to the pytest command:
+
+.. code:: shell
+
+    $ pytest tests/test-definition-ping.yaml --nornir-cache-disable
+
+

docs/source/index.rst

@@ -55,6 +55,7 @@ Contents
    Test Bundles <testbundles/alltestbundles>
    Writing Custom Tests <dev/writetests>
    Test Reports <reports/reports>
+   Advanced Options <advanced/cli>
 
 
 Indices and tables

nuts/context.py

@@ -8,10 +8,12 @@
 from nornir.core import Nornir
 from nornir.core.task import AggregatedResult, Result
 from nornir.core.filter import F
+from nornir.core.plugins.inventory import InventoryPluginRegister
 
 from nuts.helpers.errors import NutsSetupError
 from nuts.helpers.result import AbstractResultExtractor
 from nuts.helpers.filters import filter_hosts, get_filter_object
+from nuts.helpers.cache import serialize_inventory, CacheInventory
 
 
 class NutsContext:
@@ -105,17 +107,44 @@ def __init__(
         self.nornir: Optional[Nornir] = None
 
     def initialize(self) -> None:
+        """
+        Checks if inventory should be cached, then use global inventory otherwise
+        regenerate it continuously.
+        """
         if self.pytestconfig:
             config_file = pathlib.Path(
                 self.pytestconfig.getoption("nornir_configuration")
             )
         else:
             config_file = pathlib.Path(self.DEFAULT_NORNIR_CONFIG_FILE)
 
+        if self.pytestconfig and self.pytestconfig.cache:
+            if nornir_inventory := self.pytestconfig.cache.get(
+                "nuts/NORNIR_CACHE", None
+            ):
+                InventoryPluginRegister.register("NutsCacheInventory", CacheInventory)
+
+                self.nornir = InitNornir(
+                    config_file=str(config_file),
+                    logging={"enabled": False},
+                    inventory={
+                        "plugin": "NutsCacheInventory",
+                        "options": nornir_inventory,
+                    },
+                )
+                return
+
         self.nornir = InitNornir(
             config_file=str(config_file),
             logging={"enabled": False},
         )
+        if self.pytestconfig and not self.pytestconfig.getoption(
+            "nornir_cache_disabled"
+        ):
+            # pytest cash needs json encodable values
+            inventory = serialize_inventory(self.nornir.inventory)
+            if self.pytestconfig and self.pytestconfig.cache:
+                self.pytestconfig.cache.set("nuts/NORNIR_CACHE", inventory)
 
     def nuts_task(self) -> Callable[..., Result]:
         """

nuts/helpers/cache.py

@@ -0,0 +1,58 @@
+from typing import Any, Dict
+
+
+from nornir.core.inventory import Group, Groups, Host, Hosts, Inventory, ParentGroups
+from nornir.plugins.inventory.simple import _get_defaults, _get_inventory_element
+
+
+def serialize_inventory(inventory: Inventory) -> Dict[str, Dict[str, Any]]:
+    data = {
+        "hosts": {host: data.dict() for host, data in inventory.hosts.items()},
+        "groups": {group: data.dict() for group, data in inventory.groups.items()},
+        "defaults": inventory.defaults.dict(),
+    }
+    return data
+
+
+class CacheInventory:
+    def __init__(
+        self,
+        hosts: Dict[str, Dict[str, Any]],
+        groups: Dict[str, Dict[str, Any]],
+        defaults: Dict[str, Any],
+    ) -> None:
+        """
+        CacheInventory inspired by the SimpleInventory.
+
+        Args:
+
+          hosts: Dict with host name and host.dict() data
+          groups: Dict with group name and group.dict() data
+          default: defaults.dict() dict data
+        """
+
+        self.hosts_dict = hosts
+        self.groups_dict = groups
+        self.defaults_dict = defaults
+
+    def load(self) -> Inventory:
+
+        defaults = _get_defaults(self.defaults_dict)
+
+        hosts = Hosts()
+
+        for n, h in self.hosts_dict.items():
+            hosts[n] = _get_inventory_element(Host, h, n, defaults)
+
+        groups = Groups()
+
+        for n, g in self.groups_dict.items():
+            groups[n] = _get_inventory_element(Group, g, n, defaults)
+
+        for group in groups.values():
+            group.groups = ParentGroups([groups[str(g)] for g in group.groups])
+
+        for host in hosts.values():
+            host.groups = ParentGroups([groups[str(g)] for g in host.groups])
+
+        return Inventory(hosts=hosts, groups=groups, defaults=defaults)

nuts/plugin.py

@@ -103,3 +103,29 @@ def pytest_addoption(parser: Parser) -> None:
         metavar="NORNIR_CONFIG",
         help="nuts nornir configuration file. Default is nr-config.yaml",
     )
+
+    group.addoption(
+        "--nornir-cache-disable",
+        action="store_true",
+        dest="nornir_cache_disabled",
+        default=False,
+        help="disable caching of nornir inventory between executions",
+    )
+
+    group.addoption(
+        "--nornir-cached-inventory",
+        action="store_true",
+        dest="nornir_cached_inventory",
+        default=False,
+        help="Uses the chached inventory from the last executions if possible",
+    )
+
+
+def pytest_sessionstart(session: Session) -> None:
+    """Called after the ``Session`` object has been created
+    and before performing collection and entering the run test loop.
+
+    :param pytest.Session session: The pytest session object.
+    """
+    if not session.config.getoption("nornir_cached_inventory") and session.config.cache:
+        session.config.cache.set("nuts/NORNIR_CACHE", None)

tests/helpers/test_cache.py

@@ -0,0 +1,165 @@
+from nornir.core.inventory import Inventory, Hosts, Groups, Defaults, Host, Group
+from typing import Any, Dict
+
+from nuts.helpers.cache import CacheInventory, serialize_inventory
+
+
+def test_serialize_inventory():
+    hosts = Hosts(
+        {
+            "host1": Host(
+                name="host1", hostname="host1.example.com", data={"key": "value"}
+            )
+        }
+    )
+    groups = Groups({"group1": Group(name="group1", data={"key": "value"})})
+    defaults = Defaults(hostname="hostname", data={"key": "value"})
+    inventory = Inventory(hosts=hosts, groups=groups, defaults=defaults)
+
+    serialized = serialize_inventory(inventory)
+
+    expected = {
+        "hosts": {
+            "host1": {
+                "name": "host1",
+                "hostname": "host1.example.com",
+                "port": None,
+                "username": None,
+                "password": None,
+                "platform": None,
+                "data": {"key": "value"},
+                "connection_options": {},
+                "groups": [],
+            }
+        },
+        "groups": {
+            "group1": {
+                "connection_options": {},
+                "data": {"key": "value"},
+                "hostname": None,
+                "groups": [],
+                "name": "group1",
+                "password": None,
+                "platform": None,
+                "port": None,
+                "username": None,
+            }
+        },
+        "defaults": {
+            "connection_options": {},
+            "data": {"key": "value"},
+            "hostname": "hostname",
+            "password": None,
+            "platform": None,
+            "port": None,
+            "username": None,
+        },
+    }
+
+    assert serialized == expected
+
+
+def test_fail_serialize_inventory():
+    hosts = Hosts(
+        {
+            "host1": Host(
+                name="host1", hostname="host1.example.com", data={"key": "value"}
+            )
+        }
+    )
+    groups = Groups({"group1": Group(name="group1", data={"key": "value"})})
+    defaults = Defaults(hostname="hostname", data={"key": "value"})
+    inventory = Inventory(hosts=hosts, groups=groups, defaults=defaults)
+
+    serialized = serialize_inventory(inventory)
+
+    expected = {
+        "hosts": {
+            "host12": {
+                "name": "host1",
+                "hostname": "host1.example.com",
+                "port": None,
+                "username": None,
+                "password": None,
+                "platform": None,
+                "data": {"key": "value"},
+                "connection_options": {},
+                "groups": [],
+            }
+        },
+        "groups": {
+            "group12": {
+                "connection_options": {},
+                "data": {"key": "value"},
+                "hostname": None,
+                "groups": [],
+                "name": "group1",
+                "password": None,
+                "platform": None,
+                "port": None,
+                "username": None,
+            }
+        },
+        "defaults": {
+            "connection_options": {},
+            "data": {"key": "value"},
+            "hostname": "hostname",
+            "password": None,
+            "platform": None,
+            "port": None,
+            "username": "User",
+        },
+    }
+
+    assert serialized != expected
+
+
+def test_cache_inventory_load():
+    hosts_dict = {
+        "host1": {
+            "name": "host1",
+            "hostname": "host1.example.com",
+            "data": {"key": "value"},
+            "groups": [],
+            "defaults": {"key": "value"},
+        }
+    }
+    groups_dict = {
+        "group1": {
+            "name": "group1",
+            "data": {"key": "value"},
+            "groups": [],
+            "defaults": {"key": "value"},
+        }
+    }
+    defaults_dict = {"hostname": "hostname", "data": {"key": "value"}}
+
+    cache_inventory = CacheInventory(
+        hosts=hosts_dict, groups=groups_dict, defaults=defaults_dict
+    )
+
+    inventory = cache_inventory.load()
+
+    assert inventory.hosts["host1"].hostname == "host1.example.com"
+    assert inventory.hosts["host1"].data == {"key": "value"}
+    assert inventory.groups["group1"].data == {"key": "value"}
+    assert inventory.defaults.hostname == "hostname"
+    assert inventory.defaults.hostname != "not the right hostname"
+    assert inventory.defaults.data == {"key": "value"}
+
+
+def test_empty_cache_inventory_load():
+    hosts_dict: Dict[str, Dict[str, Any]] = {}
+    groups_dict: Dict[str, Dict[str, Any]] = {}
+    defaults_dict: Dict[str, Any] = {}
+
+    cache_inventory = CacheInventory(
+        hosts=hosts_dict, groups=groups_dict, defaults=defaults_dict
+    )
+
+    inventory = cache_inventory.load()
+
+    assert inventory is not None
+    assert inventory.hosts is not None
+    assert inventory.groups is not None
+    assert inventory.defaults is not None