From 58b4a873478ad8c694f25d6c1261f92b3efd16bc Mon Sep 17 00:00:00 2001
From: Lex Li <support@lextm.com>
Date: Mon, 30 Sep 2024 13:44:05 -0400
Subject: [PATCH] Added more MIB related test cases.

---
 .../asyncio/manager/cmdgen/test_v1_walk.py    | 60 +++++++++++++++++++
 ...ure-mib-viewer-and-resolve-pdu-varbinds.py | 50 ++++++++++++++++
 2 files changed, 110 insertions(+)

diff --git a/tests/hlapi/v3arch/asyncio/manager/cmdgen/test_v1_walk.py b/tests/hlapi/v3arch/asyncio/manager/cmdgen/test_v1_walk.py
index 06ae82f33..ada26dfab 100644
--- a/tests/hlapi/v3arch/asyncio/manager/cmdgen/test_v1_walk.py
+++ b/tests/hlapi/v3arch/asyncio/manager/cmdgen/test_v1_walk.py
@@ -1,5 +1,7 @@
 import pytest
 from pysnmp.hlapi.v3arch.asyncio import *
+from pysnmp.smi.builder import MibBuilder
+from pysnmp.smi.view import MibViewController
 from tests.agent_context import AGENT_PORT, AgentContextManager
 
 
@@ -36,6 +38,64 @@ async def test_v1_walk():
         snmpEngine.closeDispatcher()
 
 
+@pytest.mark.asyncio
+async def test_v1_walk_mib():
+    async with AgentContextManager():
+        mib_builder = MibBuilder()
+        mib_view_controller = MibViewController(mib_builder)
+        mib_builder.loadModules(
+            "SNMP-COMMUNITY-MIB",
+            "PYSNMP-MIB",
+            "PYSNMP-USM-MIB",
+            "SNMP-VIEW-BASED-ACM-MIB",
+        )
+
+        snmpEngine = SnmpEngine()
+        snmpEngine.cache["mibViewController"] = mib_view_controller
+        objects = walkCmd(
+            snmpEngine,
+            CommunityData("public", mpModel=0),
+            await UdpTransportTarget.create(("localhost", AGENT_PORT)),
+            ContextData(),
+            ObjectType(ObjectIdentity("SNMPv2-MIB", "sysDescr", 0)),
+        )
+
+        objects_list = [item async for item in objects]
+
+        errorIndication, errorStatus, errorIndex, varBinds = objects_list[0]
+
+        assert errorIndication is None
+        assert errorStatus == 0
+        assert len(varBinds) == 1
+        assert varBinds[0][0].prettyPrint() == "SNMPv2-MIB::sysObjectID.0"
+
+        errorIndication, errorStatus, errorIndex, varBinds = objects_list[1]
+
+        assert errorIndication is None
+        assert errorStatus == 0
+        assert len(varBinds) == 1
+        assert varBinds[0][0].prettyPrint() == "SNMPv2-MIB::sysUpTime.0"
+
+        assert len(objects_list) == 267
+
+        errorIndication, errorStatus, errorIndex, varBinds = objects_list[-1]
+        assert (
+            varBinds[0][0].prettyPrint()
+            == 'SNMP-COMMUNITY-MIB::snmpCommunityStatus."public"'
+        )
+
+        for errorIndication, errorStatus, errorIndex, varBinds in objects_list:
+            content = varBinds[0][0].prettyPrint()
+            if (
+                not content.startswith("PYSNMP-USM-MIB::")
+                and not content.startswith("SNMP-USER-BASED-SM-MIB::")
+                and not content.startswith("SNMP-VIEW-BASED-ACM-MIB::")
+            ):
+                assert content.count(".") == 1  # fully resolved.
+
+        snmpEngine.closeDispatcher()
+
+
 @pytest.mark.asyncio
 async def test_v1_walk_subtree():
     async with AgentContextManager():
diff --git a/tests/smi/manager/test_configure-mib-viewer-and-resolve-pdu-varbinds.py b/tests/smi/manager/test_configure-mib-viewer-and-resolve-pdu-varbinds.py
index cbb9ff2fb..a141579c9 100644
--- a/tests/smi/manager/test_configure-mib-viewer-and-resolve-pdu-varbinds.py
+++ b/tests/smi/manager/test_configure-mib-viewer-and-resolve-pdu-varbinds.py
@@ -48,3 +48,53 @@ def test_configure_mib_viewer_and_resolve_pdu_varbinds():
     for i, varBind in enumerate(varBinds):
         assert isinstance(varBind, rfc1902.ObjectType)
         assert varBind.prettyPrint() == expected_varbinds[i]
+
+
+def test_configure_mib_viewer_and_resolve_pdu_varbinds_full():
+    # Assemble MIB browser
+    mibBuilder = builder.MibBuilder()
+    mibViewController = view.MibViewController(mibBuilder)
+    compiler.addMibCompiler(
+        mibBuilder,
+        sources=["file:///usr/share/snmp/mibs", "https://mibs.pysnmp.com/asn1/@mib@"],
+    )
+
+    # Pre-load MIB modules we expect to work with
+    mibBuilder.loadModules("SNMPv2-MIB", "SNMP-COMMUNITY-MIB", "PYSNMP-MIB")
+
+    # Check that the MIB browser is correctly configured
+    assert isinstance(mibBuilder, builder.MibBuilder)
+    assert isinstance(mibViewController, view.MibViewController)
+
+    # Check that the expected MIB modules are loaded
+    assert mibBuilder.loadModules("SNMPv2-MIB") is not None
+    assert mibBuilder.loadModules("SNMP-COMMUNITY-MIB") is not None
+    assert mibBuilder.loadModules("PYSNMP-MIB") is not None
+
+    # Define expected var-binds
+    expected_varbinds = [
+        "SNMP-COMMUNITY-MIB::snmpTrapCommunity.0 = ",
+        "SNMPv2-MIB::snmpTrapEnterprise.0 = PYSNMP-MIB::pysnmpNotificationObjects.1.2",
+        "SNMPv2-MIB::sysDescr.0 = my system",
+    ]
+
+    # This is what we can get in TRAP PDU
+    varBinds = [
+        ("1.3.6.1.6.3.18.1.4.0", ""),
+        ("1.3.6.1.6.3.1.1.4.3.0", "1.3.6.1.4.1.20408.4.1.1.2"),
+        ("1.3.6.1.2.1.1.1.0", "my system"),
+    ]
+
+    # Run var-binds through MIB resolver
+    # You may want to catch and ignore resolution errors here
+    varBinds = [
+        rfc1902.ObjectType(rfc1902.ObjectIdentity(x[0]), x[1]).resolveWithMib(
+            mibViewController
+        )
+        for x in varBinds
+    ]
+
+    # Check that var-binds were resolved correctly
+    for i, varBind in enumerate(varBinds):
+        assert isinstance(varBind, rfc1902.ObjectType)
+        assert varBind.prettyPrint() == expected_varbinds[i]