fix test fixtures

hyperglass/execution/drivers/tests/test_construct.py

@@ -1,14 +1,28 @@
 # Project
+import typing as t
+import pytest
 from hyperglass.models.api import Query
+from hyperglass.configuration import init_ui_params
+from hyperglass.models.config.params import Params
+from hyperglass.models.directive import Directives
+from hyperglass.models.config.devices import Devices
 from hyperglass.state import use_state
-from hyperglass.test import initialize_state
 
 # Local
 from .._construct import Construct
 
+if t.TYPE_CHECKING:
+    from hyperglass.state import HyperglassState
 
-def test_construct():
-    devices = [
+
+@pytest.fixture
+def params():
+    return {}
+
+
+@pytest.fixture
+def devices():
+    return [
         {
             "name": "test1",
             "address": "127.0.0.1",
@@ -18,7 +32,11 @@ def test_construct():
             "directives": ["juniper_bgp_route"],
         }
     ]
-    directives = [
+
+
+@pytest.fixture
+def directives():
+    return [
         {
             "juniper_bgp_route": {
                 "name": "BGP Route",
@@ -27,10 +45,37 @@ def test_construct():
         }
     ]
 
-    initialize_state(params={}, directives=directives, devices=devices)
 
-    state = use_state()
+@pytest.fixture
+def state(
+    *,
+    params: t.Dict[str, t.Any],
+    directives: t.Sequence[t.Dict[str, t.Any]],
+    devices: t.Sequence[t.Dict[str, t.Any]],
+) -> t.Generator["HyperglassState", None, None]:
+    """Test fixture to initialize Redis store."""
+    _state = use_state()
+    _params = Params(**params)
+    _directives = Directives.new(*directives)
+
+    with _state.cache.pipeline() as pipeline:
+        # Write params and directives to the cache first to avoid a race condition where ui_params
+        # or devices try to access params or directives before they're available.
+        pipeline.set("params", _params)
+        pipeline.set("directives", _directives)
+
+    _devices = Devices(*devices)
+    ui_params = init_ui_params(params=_params, devices=_devices)
+
+    with _state.cache.pipeline() as pipeline:
+        pipeline.set("devices", _devices)
+        pipeline.set("ui_params", ui_params)
+
+    yield _state
+    _state.clear()
+
 
+def test_construct(state):
     query = Query(
         queryLocation="test1",
         queryTarget="192.0.2.0/24",

hyperglass/models/directive.py

@@ -7,13 +7,11 @@
 
 # Third Party
 from pydantic import (
-    Discriminator,
     field_validator,
     Field,
     FilePath,
     IPvAnyNetwork,
     PrivateAttr,
-    Tag,
 )
 
 # Project
@@ -77,7 +75,7 @@ class Select(Input):
 class Rule(HyperglassModel):
     """Base rule."""
 
-    _type: RuleTypeAttr = PrivateAttr(Field("none", discriminator="_type"))
+    _type: RuleTypeAttr = "none"
     _passed: PassedValidation = PrivateAttr(None)
     condition: Condition
     action: Action = "permit"
@@ -243,45 +241,30 @@ class RuleWithoutValidation(Rule):
     """A rule with no validation."""
 
     _type: RuleTypeAttr = "none"
-    condition: None
+    condition: None = None
 
     def validate_target(self, target: str, *, multiple: bool) -> t.Literal[True]:
         """Don't validate a target. Always returns `True`."""
         self._passed = True
         return True
 
 
-RuleWithIPv4Type = t.Annotated[RuleWithIPv4, Tag("ipv4")]
-RuleWithIPv6Type = t.Annotated[RuleWithIPv6, Tag("ipv6")]
-RuleWithPatternType = t.Annotated[RuleWithPattern, Tag("pattern")]
-RuleWithoutValidationType = t.Annotated[RuleWithoutValidation, Tag("none")]
-
-# RuleType = t.Union[RuleWithIPv4, RuleWithIPv6, RuleWithPattern, RuleWithoutValidation]
 RuleType = t.Union[
-    RuleWithIPv4Type,
-    RuleWithIPv6Type,
-    RuleWithPatternType,
-    RuleWithoutValidationType,
+    RuleWithIPv4,
+    RuleWithIPv6,
+    RuleWithPattern,
+    RuleWithoutValidation,
 ]
 
 
-def type_discriminator(value: t.Any) -> RuleTypeAttr:
-    """Pydantic type discriminator."""
-    if isinstance(value, dict):
-        return value.get("_type")
-    return getattr(value, "_type", None)
-
-
 class Directive(HyperglassUniqueModel, unique_by=("id", "table_output")):
     """A directive contains commands that can be run on a device, as long as defined rules are met."""
 
     _hyperglass_builtin: bool = PrivateAttr(False)
 
     id: str
     name: str
-    rules: t.List[RuleType] = [
-        Field(RuleWithPattern(condition="*"), discriminator=Discriminator(type_discriminator))
-    ]
+    rules: t.List[RuleType] = [RuleWithoutValidation()]
     field: t.Union[Text, Select]
     info: t.Optional[FilePath] = None
     plugins: t.List[str] = []
@@ -290,6 +273,28 @@ class Directive(HyperglassUniqueModel, unique_by=("id", "table_output")):
     multiple: bool = False
     multiple_separator: str = " "
 
+    @field_validator("rules", mode="before")
+    @classmethod
+    def validate_rules(cls, rules: t.List[t.Dict[str, t.Any]]):
+        """Initialize the correct rule type based on condition value."""
+        out_rules: t.List[RuleType] = []
+        for rule in rules:
+            if isinstance(rule, dict):
+                condition = rule.get("condition")
+                if condition is None:
+                    out_rules.append(RuleWithoutValidation(**rule))
+                try:
+                    condition_net = ip_network(condition)
+                    if condition_net.version == 4:
+                        out_rules.append(RuleWithIPv4(**rule))
+                    if condition_net.version == 6:
+                        out_rules.append(RuleWithIPv6(**rule))
+                except ValueError:
+                    out_rules.append(RuleWithPattern(**rule))
+            if isinstance(rule, Rule):
+                out_rules.append(rule)
+        return out_rules
+
     def validate_target(self, target: str) -> bool:
         """Validate a target against all configured rules."""
         for rule in self.rules:

hyperglass/plugins/tests/test_bgp_community.py

@@ -1,8 +1,17 @@
 """Test BGP Community validation."""
+import typing as t
+import pytest
 
 # Local
 from .._builtin.bgp_community import ValidateBGPCommunity
 
+from hyperglass.state import use_state
+from hyperglass.models.config.params import Params
+
+if t.TYPE_CHECKING:
+    from hyperglass.state import HyperglassState
+
+
 CHECKS = (
     ("32768", True),
     ("65000:1", True),
@@ -24,7 +33,25 @@
 )
 
 
-def test_bgp_community():
+@pytest.fixture
+def params():
+    return {}
+
+
+@pytest.fixture
+def state(*, params: t.Dict[str, t.Any]) -> t.Generator["HyperglassState", None, None]:
+    """Test fixture to initialize Redis store."""
+    _state = use_state()
+    _params = Params(**params)
+
+    with _state.cache.pipeline() as pipeline:
+        pipeline.set("params", _params)
+
+    yield _state
+    _state.clear()
+
+
+def test_bgp_community(state):
     plugin = ValidateBGPCommunity()
 
     for value, expected in CHECKS:

hyperglass/state/tests/test_hooks.py

@@ -1,9 +1,18 @@
 """Test state hooks."""
 
+import typing as t
+import pytest
+
+
+if t.TYPE_CHECKING:
+    from hyperglass.state import HyperglassState
+
 # Project
 from hyperglass.models.ui import UIParameters
-from hyperglass.models.config.params import Params
 from hyperglass.models.config.devices import Devices
+from hyperglass.configuration import init_ui_params
+from hyperglass.models.config.params import Params
+from hyperglass.models.directive import Directives
 
 # Local
 from ..hooks import use_state
@@ -13,11 +22,72 @@
     ("params", Params),
     ("devices", Devices),
     ("ui_params", UIParameters),
+    ("directives", Directives),
     (None, HyperglassState),
 )
 
 
-def test_use_state_caching():
+@pytest.fixture
+def params():
+    return {}
+
+
+@pytest.fixture
+def devices():
+    return [
+        {
+            "name": "test1",
+            "address": "127.0.0.1",
+            "credential": {"username": "", "password": ""},
+            "platform": "juniper",
+            "attrs": {"source4": "192.0.2.1", "source6": "2001:db8::1"},
+            "directives": ["juniper_bgp_route"],
+        }
+    ]
+
+
+@pytest.fixture
+def directives():
+    return [
+        {
+            "juniper_bgp_route": {
+                "name": "BGP Route",
+                "field": {"description": "test"},
+            }
+        }
+    ]
+
+
+@pytest.fixture
+def state(
+    *,
+    params: t.Dict[str, t.Any],
+    directives: t.Sequence[t.Dict[str, t.Any]],
+    devices: t.Sequence[t.Dict[str, t.Any]],
+) -> t.Generator["HyperglassState", None, None]:
+    """Test fixture to initialize Redis store."""
+    _state = use_state()
+    _params = Params(**params)
+    _directives = Directives.new(*directives)
+
+    with _state.cache.pipeline() as pipeline:
+        # Write params and directives to the cache first to avoid a race condition where ui_params
+        # or devices try to access params or directives before they're available.
+        pipeline.set("params", _params)
+        pipeline.set("directives", _directives)
+
+    _devices = Devices(*devices)
+    ui_params = init_ui_params(params=_params, devices=_devices)
+
+    with _state.cache.pipeline() as pipeline:
+        pipeline.set("devices", _devices)
+        pipeline.set("ui_params", ui_params)
+
+    yield _state
+    _state.clear()
+
+
+def test_use_state_caching(state):
     first = None
     for attr, model in STATE_ATTRS:
         for i in range(0, 5):