ci: add prepare_tests and get_coroutines benchmarks (aristanetworks#860)

* ci: add prepare_tests and get_coroutines benchmarks

* fix: clear indexes

* test: fix finventory fixture parametrization

* chore: update vscode settings

* test: fix test_runner assertion

* test: refactor benchmarks

* Update tests/benchmark/conftest.py

---------

Co-authored-by: Guillaume Mulocher <[email protected]>

.vscode/settings.json

@@ -1,15 +1,14 @@
 {
   "ruff.enable": true,
   "ruff.configuration": "pyproject.toml",
-  "python.testing.unittestEnabled": false,
   "python.testing.pytestEnabled": true,
   "python.testing.pytestArgs": [
     "tests"
   ],
-  "python.languageServer": "Pylance",
   "githubIssues.issueBranchTitle": "issues/${issueNumber}-${issueTitle}",
-  "editor.formatOnPaste": true,
-  "files.trimTrailingWhitespace": true,
-  "workbench.remoteIndicator.showExtensionRecommendations": true,
+  "pylint.importStrategy": "fromEnvironment",
+  "pylint.args": [
+    "--rcfile=pyproject.toml"
+  ],
 
 }

tests/benchmark/conftest.py

@@ -38,3 +38,20 @@ def catalog(anta_mock_env: AntaMockEnvironment) -> AntaCatalog:
 def pytest_terminal_summary(terminalreporter: TerminalReporter) -> None:
     """Display the total number of ANTA unit test cases used to benchmark."""
     terminalreporter.write_sep("=", f"{TEST_CASE_COUNT} ANTA test cases")
+
+
+def pytest_generate_tests(metafunc: pytest.Metafunc) -> None:
+    """Parametrize inventory for benchmark tests."""
+    if "inventory" in metafunc.fixturenames:
+        for marker in metafunc.definition.iter_markers(name="parametrize"):
+            if "inventory" in marker.args[0]:
+                # Do not override test function parametrize marker for inventory arg
+                return
+        metafunc.parametrize(
+            "inventory",
+            [
+                pytest.param({"count": 1, "disable_cache": True, "reachable": True}, id="1-device"),
+                pytest.param({"count": 2, "disable_cache": True, "reachable": True}, id="2-devices"),
+            ],
+            indirect=True,
+        )

tests/benchmark/test_anta.py

@@ -22,27 +22,18 @@
 logger = logging.getLogger(__name__)
 
 
-@pytest.mark.parametrize(
-    "inventory",
-    [
-        pytest.param({"count": 1, "disable_cache": True, "reachable": False}, id="1 device"),
-        pytest.param({"count": 2, "disable_cache": True, "reachable": False}, id="2 devices"),
-    ],
-    indirect=True,
-)
-def test_anta_dry_run(benchmark: BenchmarkFixture, catalog: AntaCatalog, inventory: AntaInventory) -> None:
-    """Test and benchmark ANTA in Dry-Run Mode."""
+def test_anta_dry_run(benchmark: BenchmarkFixture, event_loop: asyncio.AbstractEventLoop, catalog: AntaCatalog, inventory: AntaInventory) -> None:
+    """Benchmark ANTA in Dry-Run Mode."""
     # Disable logging during ANTA execution to avoid having these function time in benchmarks
     logging.disable()
 
-    def bench() -> ResultManager:
-        """Need to wrap the ANTA Runner to instantiate a new ResultManger for each benchmark run."""
+    def _() -> ResultManager:
         manager = ResultManager()
         catalog.clear_indexes()
-        asyncio.run(main(manager, inventory, catalog, dry_run=True))
+        event_loop.run_until_complete(main(manager, inventory, catalog, dry_run=True))
         return manager
 
-    manager = benchmark(bench)
+    manager = benchmark(_)
 
     logging.disable(logging.NOTSET)
     if len(manager.results) != len(inventory) * len(catalog.tests):
@@ -51,30 +42,21 @@ def bench() -> ResultManager:
     logger.info(bench_info)
 
 
-@pytest.mark.parametrize(
-    "inventory",
-    [
-        pytest.param({"count": 1, "disable_cache": True}, id="1 device"),
-        pytest.param({"count": 2, "disable_cache": True}, id="2 devices"),
-    ],
-    indirect=True,
-)
 @patch("anta.models.AntaTest.collect", collect)
 @patch("anta.device.AntaDevice.collect_commands", collect_commands)
 @respx.mock  # Mock eAPI responses
-def test_anta(benchmark: BenchmarkFixture, catalog: AntaCatalog, inventory: AntaInventory) -> None:
-    """Test and benchmark ANTA. Mock eAPI responses."""
+def test_anta(benchmark: BenchmarkFixture, event_loop: asyncio.AbstractEventLoop, catalog: AntaCatalog, inventory: AntaInventory) -> None:
+    """Benchmark ANTA."""
     # Disable logging during ANTA execution to avoid having these function time in benchmarks
     logging.disable()
 
-    def bench() -> ResultManager:
-        """Need to wrap the ANTA Runner to instantiate a new ResultManger for each benchmark run."""
+    def _() -> ResultManager:
         manager = ResultManager()
         catalog.clear_indexes()
-        asyncio.run(main(manager, inventory, catalog))
+        event_loop.run_until_complete(main(manager, inventory, catalog))
         return manager
 
-    manager = benchmark(bench)
+    manager = benchmark(_)
 
     logging.disable(logging.NOTSET)
 

tests/benchmark/test_runner.py

@@ -0,0 +1,48 @@
+# Copyright (c) 2023-2024 Arista Networks, Inc.
+# Use of this source code is governed by the Apache License 2.0
+# that can be found in the LICENSE file.
+"""Benchmark tests for anta.runner."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from anta.result_manager import ResultManager
+from anta.runner import get_coroutines, prepare_tests
+
+if TYPE_CHECKING:
+    from collections import defaultdict
+
+    from pytest_codspeed import BenchmarkFixture
+
+    from anta.catalog import AntaCatalog, AntaTestDefinition
+    from anta.device import AntaDevice
+    from anta.inventory import AntaInventory
+
+
+def test_prepare_tests(benchmark: BenchmarkFixture, catalog: AntaCatalog, inventory: AntaInventory) -> None:
+    """Benchmark `anta.runner.prepare_tests`."""
+
+    def _() -> defaultdict[AntaDevice, set[AntaTestDefinition]] | None:
+        catalog.clear_indexes()
+        return prepare_tests(inventory=inventory, catalog=catalog, tests=None, tags=None)
+
+    selected_tests = benchmark(_)
+
+    assert selected_tests is not None
+    assert len(selected_tests) == len(inventory)
+    assert sum(len(tests) for tests in selected_tests.values()) == len(inventory) * len(catalog.tests)
+
+
+def test_get_coroutines(benchmark: BenchmarkFixture, catalog: AntaCatalog, inventory: AntaInventory) -> None:
+    """Benchmark `anta.runner.get_coroutines`."""
+    selected_tests = prepare_tests(inventory=inventory, catalog=catalog, tests=None, tags=None)
+
+    assert selected_tests is not None
+
+    coroutines = benchmark(lambda: get_coroutines(selected_tests=selected_tests, manager=ResultManager()))
+    for coros in coroutines:
+        coros.close()
+
+    count = sum(len(tests) for tests in selected_tests.values())
+    assert count == len(coroutines)

tests/conftest.py

@@ -17,18 +17,20 @@
 DATA_DIR: Path = Path(__file__).parent.resolve() / "data"
 
 
-@pytest.fixture(params=[{"count": 1}], ids=["1-reachable-device-without-cache"])
+@pytest.fixture
 def inventory(request: pytest.FixtureRequest) -> Iterator[AntaInventory]:
     """Generate an ANTA inventory."""
     user = "admin"
     password = "password"  # noqa: S105
-    disable_cache = request.param.get("disable_cache", True)
-    reachable = request.param.get("reachable", True)
-    if "filename" in request.param:
-        inv = AntaInventory.parse(DATA_DIR / request.param["filename"], username=user, password=password, disable_cache=disable_cache)
+    params = request.param if hasattr(request, "param") else {}
+    count = params.get("count", 1)
+    disable_cache = params.get("disable_cache", True)
+    reachable = params.get("reachable", True)
+    if "filename" in params:
+        inv = AntaInventory.parse(DATA_DIR / params["filename"], username=user, password=password, disable_cache=disable_cache)
     else:
         inv = AntaInventory()
-        for i in range(request.param["count"]):
+        for i in range(count):
             inv.add_device(
                 AsyncEOSDevice(
                     host=f"device-{i}.anta.arista.com",