feat: add --learn mode for expected state capture

nac_test/cli/main.py

@@ -279,6 +279,16 @@ def version_callback(value: bool) -> None:
 ]
 
 
+Learn = Annotated[
+    bool,
+    typer.Option(
+        "--learn",
+        help="Run in learning mode: capture live operational state as baseline instead of verifying. Output is written to {output}/learned_state/ and can be loaded as a -d path for verification.",
+        envvar="NAC_TEST_LEARN",
+    ),
+]
+
+
 Testbed = Annotated[
     Path | None,
     typer.Option(
@@ -304,6 +314,7 @@ def main(
     exclude: Exclude = None,
     render_only: RenderOnly = False,
     dry_run: DryRun = False,
+    learn: Learn = False,
     processes: Processes = None,
     pyats: PyATS = False,
     robot: Robot = False,
@@ -413,6 +424,7 @@ def main(
         exclude_tags=exclude,
         render_only=render_only,
         dry_run=dry_run,
+        learn=learn,
         processes=processes,
         extra_args=validated_robot_args,
         max_parallel_devices=max_parallel_devices,

nac_test/combined_orchestrator.py

@@ -82,6 +82,7 @@ def __init__(
         exclude_tags: list[str] | None = None,
         render_only: bool = False,
         dry_run: bool = False,
+        learn: bool = False,
         max_parallel_devices: int | None = None,
         minimal_reports: bool = False,
         custom_testbed_path: Path | None = None,
@@ -129,6 +130,7 @@ def __init__(
         self.exclude_tags = exclude_tags or []
         self.render_only = render_only
         self.dry_run = dry_run
+        self.learn = learn
         self.processes = processes
         self.extra_args = extra_args
 
@@ -220,6 +222,7 @@ def run_tests(self) -> CombinedResults:
                 custom_testbed_path=self.custom_testbed_path,
                 controller_type=self.controller_type,
                 dry_run=self.dry_run,
+                learn=self.learn,
                 verbose=self.verbose,
                 loglevel=self.loglevel,
                 include_tags=self.include_tags,

nac_test/pyats_core/common/base_test.py

@@ -50,6 +50,7 @@
     get_defaults_prefix,
 )
 from nac_test.utils.formatting import format_file_timestamp_ms
+from nac_test.utils.learned_state import save_learned_state
 from nac_test.utils.yaml import safe_load
 
 T = TypeVar("T")
@@ -1956,6 +1957,12 @@ async def run_verification_async(self) -> list[VerificationResult]:
                 }
             ]
 
+        # Learning mode: capture state instead of verifying
+        if getattr(self, "SUPPORTS_LEARNING", False) and os.environ.get(
+            "NAC_TEST_LEARN"
+        ):
+            return await self._run_learning_mode(items_to_verify)
+
         # Detect verification pattern based on return type
         if isinstance(items_to_verify, dict):
             # Grouped verification: {group_key: [contexts]}
@@ -2068,6 +2075,92 @@ async def _run_grouped_verification(
 
         return flattened_results
 
+    async def _run_learning_mode(
+        self, items: list[dict[str, Any]] | dict[str, list[dict[str, Any]]]
+    ) -> list["VerificationResult"]:
+        """Execute in learning mode — capture state instead of asserting.
+
+        Calls the test's capture_learned_state() method to query live state,
+        then writes the captured data to the learned_state output directory.
+
+        Args:
+            items: Items from get_items_to_verify() (list or dict depending on pattern).
+
+        Returns:
+            Single-element list with a PASSED result indicating successful capture.
+        """
+        # Flatten grouped items to a list for the capture method
+        if isinstance(items, dict):
+            flat_items = [ctx for group in items.values() for ctx in group]
+        else:
+            flat_items = items
+
+        from nac_test.pyats_core.constants import DEFAULT_API_CONCURRENCY
+
+        semaphore = asyncio.Semaphore(DEFAULT_API_CONCURRENCY)
+        client = getattr(self, "client", None)
+
+        if not hasattr(self, "capture_learned_state"):
+            msg = (
+                f"{self.__class__.__name__} has SUPPORTS_LEARNING=True but no "
+                f"capture_learned_state() method. Inherit LearningModeMixin."
+            )
+            self.logger.warning(msg)
+            return [
+                {
+                    "status": ResultStatus.SKIPPED,
+                    "context": {"action": "learn"},
+                    "reason": msg,
+                    "api_duration": 0,
+                }
+            ]
+
+        try:
+            learned_data = await self.capture_learned_state(
+                semaphore, client, flat_items
+            )
+        except NotImplementedError as e:
+            self.logger.warning(str(e))
+            return [
+                {
+                    "status": ResultStatus.SKIPPED,
+                    "context": {"action": "learn"},
+                    "reason": str(e),
+                    "api_duration": 0,
+                }
+            ]
+
+        # Write captured state to output directory
+        hostname = getattr(self, "hostname", None)
+        test_name = self.__class__.__name__
+        learned_state_dir = Path(
+            os.environ.get("NAC_TEST_LEARNED_STATE_DIR", "learned_state")
+        )
+
+        # If capture returned empty data, write a marker file with the reason
+        # so users know learning ran but found nothing (vs. never ran at all)
+        if not learned_data:
+            config = getattr(self, "TEST_CONFIG", {})
+            command = config.get("api_endpoint", "unknown")
+            learned_data = {
+                "_learned_state_empty": f"No data returned by parser for: {command}"
+            }
+
+        output_path = save_learned_state(
+            learned_data, learned_state_dir, test_name, hostname
+        )
+
+        self.logger.info(f"Learned state saved to {output_path}")
+
+        return [
+            {
+                "status": ResultStatus.PASSED,
+                "context": {"action": "learn", "output_path": str(output_path)},
+                "reason": f"Successfully captured baseline state ({len(flat_items)} items) → {output_path.name}",
+                "api_duration": 0,
+            }
+        ]
+
     async def _run_item_verification(
         self, items: list[dict[str, Any]]
     ) -> list[VerificationResult]:

nac_test/pyats_core/common/learning_mode_mixin.py

@@ -0,0 +1,119 @@
+# SPDX-License-Identifier: MPL-2.0
+# Copyright (c) 2025 Daniel Schmidt
+
+"""Learning mode mixin for operational test classes.
+
+Provides the contract for tests that support a two-phase workflow:
+1. Learn mode (--learn): capture live state and save as baseline
+2. Verify mode (default): compare live state against captured baseline
+
+Tests opt in by inheriting LearningModeMixin and implementing
+capture_learned_state(). The base class orchestration checks for
+SUPPORTS_LEARNING and the NAC_TEST_LEARN env var to route execution.
+
+Usage:
+    class VerifyBGPRoutes(LearningModeMixin, IOSXETestBase):
+        SUPPORTS_LEARNING = True
+
+        async def capture_learned_state(self, semaphore, client, items):
+            # Query live state and return structured data
+            return {"sdwan": {"sites": [...]}}
+
+        def get_items_to_verify(self):
+            # Same as normal — extract what to check
+            ...
+
+        async def verify_item(self, semaphore, client, context):
+            # Normal verification against data model (which now includes learned state)
+            ...
+"""
+
+import os
+from pathlib import Path
+from typing import Any
+
+from pyats import aetest
+
+
+class LearningModeMixin(aetest.Testcase):  # type: ignore[misc]
+    """Mixin adding learning mode support to operational test classes.
+
+    Inherits from aetest.Testcase so that PyATS's TestableMeta metaclass
+    processes this class correctly (methods get the required .source attribute).
+    Python's MRO ensures aetest.Testcase appears only once when combined with
+    other base classes that also inherit from it.
+
+    Tests that support learning inherit this mixin and override
+    capture_learned_state(). The framework detects learn mode via
+    the NAC_TEST_LEARN environment variable and calls the capture
+    method instead of the normal verify loop.
+
+    Usage:
+        class MyTest(LearningModeMixin, SDWANManagerTestBase):
+            async def capture_learned_state(self, semaphore, client, items):
+                ...
+
+    Attributes:
+        SUPPORTS_LEARNING: Class-level flag indicating this test supports
+            the --learn mode. Set to True in subclasses that implement
+            capture_learned_state().
+        LEARNED_STATE_KEY: Unique string identifying this test's namespace
+            within the learned_state data structure. Must be unique across
+            all tests that support learning. Used as the key under
+            {architecture}.learned_state.{LEARNED_STATE_KEY} in the merged
+            data model. Convention: use the class name.
+    """
+
+    SUPPORTS_LEARNING: bool = True
+    LEARNED_STATE_KEY: str = ""
+
+    @property
+    def is_learn_mode(self) -> bool:
+        """Check if running in learning mode.
+
+        Returns:
+            True if NAC_TEST_LEARN environment variable is set and truthy.
+        """
+        return bool(os.environ.get("NAC_TEST_LEARN"))
+
+    @property
+    def learned_state_dir(self) -> Path:
+        """Get the output directory for learned state files.
+
+        Returns:
+            Path from NAC_TEST_LEARNED_STATE_DIR env var, or 'learned_state'
+            as fallback.
+        """
+        return Path(os.environ.get("NAC_TEST_LEARNED_STATE_DIR", "learned_state"))
+
+    async def capture_learned_state(
+        self,
+        semaphore: Any,
+        client: Any,
+        items: list[dict[str, Any]],
+    ) -> dict[str, Any]:
+        """Capture live state for all items. Override in subclass.
+
+        Called in learning mode instead of the normal verify_item() loop.
+        The implementation should make the same queries as verify_item() but
+        return the raw captured state rather than a pass/fail verdict.
+
+        The returned dict should be structured so it merges cleanly into the
+        data model when loaded via -d (using nac_yaml's merge_dict logic).
+
+        Args:
+            semaphore: Asyncio semaphore for concurrency control.
+            client: HTTP client or SSH connection (same as verify_item receives).
+            items: List of context dicts from get_items_to_verify().
+
+        Returns:
+            Dictionary containing captured state, structured for data model merge.
+
+        Raises:
+            NotImplementedError: If subclass doesn't override this method.
+        """
+        raise NotImplementedError(
+            f"{self.__class__.__name__} has SUPPORTS_LEARNING=True but does not "
+            f"implement capture_learned_state(). Override this method to define "
+            f"what state to capture in learning mode."
+        )

nac_test/pyats_core/orchestrator.py

@@ -71,6 +71,7 @@ def __init__(
         custom_testbed_path: Path | None = None,
         controller_type: str | None = None,
         dry_run: bool = False,
+        learn: bool = False,
         verbose: bool = False,
         loglevel: LogLevel = DEFAULT_LOGLEVEL,
         include_tags: list[str] | None = None,
@@ -108,6 +109,7 @@ def __init__(
         self.minimal_reports = minimal_reports
         self.custom_testbed_path = custom_testbed_path
         self.dry_run = dry_run
+        self.learn = learn
         self.verbose = verbose
         self.loglevel = loglevel
         self.include_tags = include_tags
@@ -278,6 +280,12 @@ async def _execute_api_tests_standard(self, test_files: list[Path]) -> Path | No
             env["NAC_TEST_TYPE"] = "api"
             # Pass test_dir so plugin can compute relative test names
             env[ENV_TEST_DIR] = str(self.test_dir)
+            # Learning mode: propagate flag and set output directory
+            if self.learn:
+                env["NAC_TEST_LEARN"] = "1"
+                learned_state_dir = Path.cwd() / "learned_state"
+                learned_state_dir.mkdir(parents=True, exist_ok=True)
+                env["NAC_TEST_LEARNED_STATE_DIR"] = str(learned_state_dir)
 
             # Execute and return the archive path
             assert self.subprocess_runner is not None  # Should be initialized by now
@@ -365,6 +373,13 @@ async def _execute_ssh_tests_device_centric(
                 # Set environment variable for test subprocesses to find broker
                 os.environ["NAC_TEST_BROKER_SOCKET"] = str(broker.socket_path)
 
+                # Learning mode: propagate to D2D subprocesses via os.environ
+                if self.learn:
+                    os.environ["NAC_TEST_LEARN"] = "1"
+                    learned_state_dir = Path.cwd() / "learned_state"
+                    learned_state_dir.mkdir(parents=True, exist_ok=True)
+                    os.environ["NAC_TEST_LEARNED_STATE_DIR"] = str(learned_state_dir)
+
                 # Execute device tests with broker running
                 return await self._execute_device_tests_with_broker(test_files, devices)
 
@@ -374,8 +389,10 @@ async def _execute_ssh_tests_device_centric(
             )
             return None
         finally:
-            # Clean up environment variable
+            # Clean up environment variables
             os.environ.pop("NAC_TEST_BROKER_SOCKET", None)
+            os.environ.pop("NAC_TEST_LEARN", None)
+            os.environ.pop("NAC_TEST_LEARNED_STATE_DIR", None)
 
     async def _execute_device_tests_with_broker(
         self, test_files: list[Path], devices: list[dict[str, Any]]