Make TestNapalmBgpNeighbors and TestNapalmBgpNeighborsCount vrf aware

docs/source/testbundles/alltestbundles.rst

@@ -99,6 +99,8 @@ BGP Neighbors - Information
 .. code:: yaml
 
     - test_class: TestNapalmBgpNeighbors
+      test_execution: 
+        vrf: <string, optional>
       test_data:
         - host: <host name, required>
           local_id: <ID>
@@ -118,6 +120,9 @@ Required fields for specific tests in this bundle:
     * Test if the peer is enabled: ``host, peer, is_enabled``
     * Test if the peer is up: ``host, peer, is_up``
 
+``test_execution`` ``vrf``: This field is used to select the scope. Default is the "global" VRF. 
+
+
 **Test Bundle Example:**
 
 .. code:: yaml
@@ -147,10 +152,13 @@ BGP Neighbors - Count
 .. code:: yaml
 
     - test_class: TestNapalmBgpNeighborsCount
+      test_execution: 
+        vrf: <string, optional>
       test_data:
         - host: <host name, required>
           neighbor_count: <number of neighbors, required>
 
+``test_execution`` ``vrf``: This field is used to select the scope. Default is the "global" VRF. 
 
 **Test Bundle Example:**
 

nuts/base_tests/napalm_bgp_neighbors.py

@@ -1,6 +1,6 @@
 """Query BGP neighbors of a device or count them."""
 
-from typing import Dict, Callable, List, Any
+from typing import Dict, Callable, List, Any, Tuple
 
 import pytest
 from nornir.core.task import MultiResult, Result
@@ -13,15 +13,23 @@
 class BgpNeighborsExtractor(AbstractHostResultExtractor):
     def single_transform(
         self, single_result: MultiResult
-    ) -> Dict[str, Dict[str, NutsResult]]:
+    ) -> Dict[Tuple[str, str], Dict[str, NutsResult]]:
+        test_exectuion = self._nuts_ctx.nuts_parameters.get("test_execution")
+        if test_exectuion:
+            vrf = test_exectuion.get("vrf", "global")
+            if not vrf:
+                vrf = "global"
+        else:
+            vrf = "global"
+
         neighbors = self._simple_extract(single_result)["bgp_neighbors"]
-        if "global" not in neighbors:
+        if vrf not in neighbors:
             return {}
-        global_scope = neighbors["global"]
-        router_id = global_scope["router_id"]
+        scope = neighbors[vrf]
+        router_id = scope["router_id"]
         return {
             peer: self._add_local_id(details, router_id)
-            for peer, details in global_scope["peers"].items()
+            for peer, details in scope["peers"].items()
         }
 
     def _add_local_id(self, element: Dict[str, Any], router_id: str) -> Dict[str, Any]:
@@ -34,6 +42,8 @@ def nuts_task(self) -> Callable[..., Result]:
         return napalm_get
 
     def nuts_arguments(self) -> Dict[str, List[str]]:
+        # Overrides nuts_arguments.nuts_arguments
+        # "test_execution" is not passed to the nuts_task
         return {"getters": ["bgp_neighbors"]}
 
     def nuts_extractor(self) -> BgpNeighborsExtractor:
@@ -49,7 +59,9 @@ def test_neighbor_count(
         self, single_result: NutsResult, neighbor_count: int
     ) -> None:
         assert single_result.result is not None
-        assert len(single_result.result) == neighbor_count
+        assert (
+            len(single_result.result) == neighbor_count
+        ), f"Expected {neighbor_count}; got {len(single_result.result)}"
 
 
 class TestNapalmBgpNeighbors:

tests/base_tests/conftest.py

@@ -114,13 +114,23 @@ def transformed_result(
 
 @pytest.fixture
 def integration_tester(monkeypatch, pytester, default_nr_init):
-    def _run(selftestdata, *, test_class, task_module, task_name, test_count):
+    def _run(
+        selftestdata,
+        *,
+        test_class,
+        task_module,
+        task_name,
+        test_count,
+        test_execution=None,
+    ):
         yaml_data = [
             {
                 "test_class": test_class,
                 "test_data": [selftestdata.test_data],
             }
         ]
+        if test_execution:
+            yaml_data[0]["test_execution"] = test_execution
 
         yaml_file = pytester.path / f"test{YAML_EXTENSION}"
         with yaml_file.open("w") as fo:

tests/base_tests/test_napalm_bgp_neighbors.py

@@ -84,6 +84,13 @@
     }
 }
 
+nornir_raw_r1_vrf = {
+    "bgp_neighbors": {"CustA": nornir_raw_r1["bgp_neighbors"]["global"]}
+}
+
+nornir_raw_r2_vrf = {
+    "bgp_neighbors": {"CustA": nornir_raw_r2["bgp_neighbors"]["global"]}
+}
 
 bgp_r1_1 = SelfTestData(
     name="r1_1",
@@ -130,6 +137,21 @@
     },
 )
 
+bgp_r2_vrf = SelfTestData(
+    name="r2",
+    nornir_raw_result=nornir_raw_r2_vrf,
+    test_data={
+        "host": "R2",
+        "local_id": R2_IP,
+        "peer": R1_IP,
+        "is_enabled": True,
+        "remote_as": AS1,
+        "remote_id": AS_ID,
+        "local_as": AS2,
+        "is_up": False,
+    },
+)
+
 bgp_r1_count = SelfTestData(
     name="r1_count",
     nornir_raw_result=nornir_raw_r1,
@@ -148,6 +170,24 @@
     },
 )
 
+bgp_r1_count_vrf = SelfTestData(
+    name="r1_count",
+    nornir_raw_result=nornir_raw_r1_vrf,
+    test_data={
+        "host": "R1",
+        "neighbor_count": 2,
+    },
+)
+
+bgp_r2_count_vrf = SelfTestData(
+    name="r2_count",
+    nornir_raw_result=nornir_raw_r2_vrf,
+    test_data={
+        "host": "R2",
+        "neighbor_count": 1,
+    },
+)
+
 
 @pytest.fixture
 def general_result(timeouted_multiresult):
@@ -176,6 +216,16 @@ def selftestdata(request):
     return request.param
 
 
+@pytest.fixture(
+    params=[
+        bgp_r2_vrf,
+    ],
+    ids=lambda data: data.name,
+)
+def selftestdata_vrf(request):
+    return request.param
+
+
 @pytest.fixture
 def testdata(selftestdata):
     return selftestdata.test_data
@@ -189,6 +239,14 @@ def selftestdata_countneighbors(request):
     return request.param
 
 
+@pytest.fixture(
+    params=[bgp_r1_count_vrf, bgp_r2_count_vrf],
+    ids=lambda data: data.name,
+)
+def selftestdata_countneighbors_vrf(request):
+    return request.param
+
+
 pytestmark = [pytest.mark.nuts_test_ctx(CONTEXT())]
 
 
@@ -249,6 +307,17 @@ def test_integration(selftestdata, integration_tester):
     )
 
 
+def test_integration_vrf(selftestdata_vrf, integration_tester):
+    integration_tester(
+        selftestdata_vrf,
+        test_class="TestNapalmBgpNeighbors",
+        task_module=tasks,
+        test_execution={"vrf": "CustA"},
+        task_name="napalm_get",
+        test_count=6,
+    )
+
+
 def test_integration_count(selftestdata_countneighbors, integration_tester):
     integration_tester(
         selftestdata_countneighbors,
@@ -257,3 +326,14 @@ def test_integration_count(selftestdata_countneighbors, integration_tester):
         task_name="napalm_get",
         test_count=1,
     )
+
+
+def test_integration_count_vrf(selftestdata_countneighbors_vrf, integration_tester):
+    integration_tester(
+        selftestdata_countneighbors_vrf,
+        test_class="TestNapalmBgpNeighborsCount",
+        task_module=tasks,
+        test_execution={"vrf": "CustA"},
+        task_name="napalm_get",
+        test_count=1,
+    )