test: Enable tests from fixtures to work with string-valued Enum params (#1841)

Author: morrison12

tests/operations/selinux.boolean/change_bool_name_with_metachars.json

@@ -1,14 +1,6 @@
 {
-    "args": ["httpd`id`", "on"],
-    "kwargs": {
-        "persistent": true
-    },
-    "facts": {
-        "selinux.SEBoolean": {
-            "boolean=httpd`id`": "off"
-        }
-    },
-    "commands": [
-        "setsebool -P 'httpd`id`' on"
-    ]
+    "args": ["httpd`id`", "enum:Boolean:on"],
+    "kwargs": {"persistent": true},
+    "facts": {"selinux.SEBoolean": {"boolean=httpd`id`": "off"}},
+    "commands": ["setsebool -P 'httpd`id`' on"]
 }

tests/operations/selinux.boolean/change_not_persistent.json

@@ -1,13 +1,6 @@
 {
-    "args": ["some_boolean", "off"],
-    "kwargs": {
-    },
-    "facts": {
-        "selinux.SEBoolean": {
-            "boolean=some_boolean": "on"
-        }
-    },
-    "commands": [
-        "setsebool some_boolean off"
-    ]
+    "args": ["some_boolean", "enum:Boolean:off"],
+    "kwargs": {},
+    "facts": {"selinux.SEBoolean": {"boolean=some_boolean": "on"}},
+    "commands": ["setsebool some_boolean off"]
 }

tests/operations/selinux.boolean/change_persistent.json

@@ -1,14 +1,6 @@
 {
-    "args": ["some_boolean", "on"],
-    "kwargs": {
-        "persistent":true
-    },
-    "facts": {
-        "selinux.SEBoolean": {
-            "boolean=some_boolean": "off"
-        }
-    },
-    "commands": [
-        "setsebool -P some_boolean on"
-    ]
+    "args": ["some_boolean", "enum:Boolean:on"],
+    "kwargs": {"persistent": true},
+    "facts": {"selinux.SEBoolean": {"boolean=some_boolean": "off"}},
+    "commands": ["setsebool -P some_boolean on"],
 }

tests/operations/selinux.boolean/no_change.json

@@ -1,14 +1,7 @@
 {
-    "args": ["some_boolean", "on"],
-    "kwargs": {
-        "persistent":true
-    },
-    "facts": {
-        "selinux.SEBoolean": {
-            "boolean=some_boolean": "on"
-        }
-    },
-    "commands": [
-    ],
+    "args": ["some_boolean", "enum:Boolean:on"],
+    "kwargs": {"persistent": true},
+    "facts": {"selinux.SEBoolean": {"boolean=some_boolean": "on"}},
+    "commands": [],
     "noop_description": "boolean 'some_boolean' already had the value 'on'"
 }

tests/test_facts.py

@@ -12,7 +12,7 @@
 from pyinfra.context import ctx_host, ctx_state
 from pyinfra_cli.util import json_encode
 
-from .util import FakeState, create_host, get_command_string
+from .util import FakeState, create_host, get_command_string, get_enum_map, parse_value
 
 # show full diff on json
 TestCase.maxDiff = None
@@ -60,7 +60,7 @@ def _test_fn(self, test_name, test_data, fact):
                 short_fact = fact
                 fact = fact.fact()
 
-            test_args = test_data.get("arg", [])
+            test_args = parse_value(test_data.get("arg", []), get_enum_map(fact.command))
             command = _make_command(fact.command, test_args)
 
             if "command" in test_data:

tests/test_operations.py

@@ -12,7 +12,7 @@
 from pyinfra.context import ctx_host, ctx_state
 from pyinfra_cli.util import json_encode
 
-from .util import FakeState, create_host, get_command_string, parse_value, patch_files
+from .util import FakeState, create_host, get_command_string, get_enum_map, parse_value, patch_files
 
 PLATFORM_NAME = platform.system()
 
@@ -116,13 +116,15 @@ def _test(self, test_name, test_data):
 
             op_test_name = f"{arg}/{test_name}.json"
 
-            # Create a host with this tests facts and attach to context host
-            host = create_host(self.state, facts=test_data.get("facts", {}))
+            enum_map = get_enum_map(op._inner)  # noqa: SLF001
+
+            # Create a host with this test's facts and attach to context host
+            host = create_host(self.state, facts=parse_value(test_data.get("facts", {}), enum_map))
 
             allowed_exception = test_data.get("exception")
 
-            args = parse_value(test_data.get("args", []))
-            kwargs = parse_value(test_data.get("kwargs", {}))
+            args = parse_value(test_data.get("args", []), enum_map)
+            kwargs = parse_value(test_data.get("kwargs", {}), enum_map)
 
             with ctx_state.use(self.state):
                 with ctx_host.use(host):

tests/test_operations_utils.py

@@ -1,3 +1,4 @@
+from enum import Enum
 from unittest import TestCase
 from unittest.mock import MagicMock
 
@@ -8,6 +9,8 @@
 from pyinfra.operations.util.files import ensure_mode_int, unix_path_join
 from pyinfra.operations.util.packaging import ensure_packages
 
+from .util import get_enum_map
+
 
 class TestUnixPathJoin(TestCase):
     def test_simple_path(self):
@@ -545,3 +548,93 @@ def test_new_format_multi_version_pinned_match_against_any_version(self):
         )
         assert commands == []
         host.noop.assert_called_once_with("package kernel is installed (5.10.0-26,6.1.0-13)")
+
+
+class TestGetEnumMap(TestCase):
+    """ensure get_enum_map works correctly"""
+
+    class ZipEnum(Enum):
+        FOO = 1
+        BAR = 2
+
+    class ZapEnum(Enum):
+        FOO = 1
+        BAR = 2
+
+    class ZowieEnum(Enum):
+        FOO = 1
+        BAR = 2
+
+    def test_multiple_simple_types_no_enum_works(self):
+        def func(a: int, b: str, c: float, d: complex) -> None:
+            pass
+
+        assert get_enum_map(func) == {}
+
+    def test_only_simple_enum_works(self) -> None:
+        def func(a: TestGetEnumMap.ZipEnum) -> None:
+            pass
+
+        assert get_enum_map(func) == {"ZipEnum": TestGetEnumMap.ZipEnum}
+
+    def test_union_simple_enum_or_none_works(self) -> None:
+        def func(a: TestGetEnumMap.ZipEnum | None) -> None:
+            pass
+
+        assert get_enum_map(func) == {"ZipEnum": TestGetEnumMap.ZipEnum}
+
+    def test_union_simple_enum_or_str_works(self) -> None:
+        def func(a: TestGetEnumMap.ZipEnum | str) -> None:
+            pass
+
+        assert get_enum_map(func) == {"ZipEnum": TestGetEnumMap.ZipEnum}
+
+    def test_list_of_simple_enum_works(self) -> None:
+        def func(a: list[TestGetEnumMap.ZipEnum]) -> None:
+            pass
+
+        assert get_enum_map(func) == {"ZipEnum": TestGetEnumMap.ZipEnum}
+
+    def test_set_of_simple_enum_works(self) -> None:
+        def func(a: set[TestGetEnumMap.ZipEnum]) -> None:
+            pass
+
+        assert get_enum_map(func) == {"ZipEnum": TestGetEnumMap.ZipEnum}
+
+    def test_tuple_of_simple_enum_works(self) -> None:
+        def func(a: tuple[TestGetEnumMap.ZipEnum, ...]) -> None:
+            pass
+
+        assert get_enum_map(func) == {"ZipEnum": TestGetEnumMap.ZipEnum}
+
+    def test_dict_key_of_simple_enum_works(self) -> None:
+        def func(a: dict[TestGetEnumMap.ZipEnum, str]) -> None:
+            pass
+
+        assert get_enum_map(func) == {"ZipEnum": TestGetEnumMap.ZipEnum}
+
+    def test_dict_value_of_simple_enum_works(self) -> None:
+        def func(a: dict[str, TestGetEnumMap.ZipEnum]) -> None:
+            pass
+
+        assert get_enum_map(func) == {"ZipEnum": TestGetEnumMap.ZipEnum}
+
+    def test_multiple_enums_in_unions_work(self) -> None:
+        def func(
+            a: TestGetEnumMap.ZipEnum | None,
+            b: TestGetEnumMap.ZapEnum | str,
+            c: TestGetEnumMap.ZowieEnum | complex,
+        ) -> None:
+            pass
+
+        assert get_enum_map(func) == {
+            "ZipEnum": TestGetEnumMap.ZipEnum,
+            "ZapEnum": TestGetEnumMap.ZapEnum,
+            "ZowieEnum": TestGetEnumMap.ZowieEnum,
+        }
+
+    def test_nested_list_of_enum_unions_works(self) -> None:
+        def func(a: list[list[TestGetEnumMap.ZapEnum | None] | None] | None) -> None:
+            pass
+
+        assert get_enum_map(func) == {"ZapEnum": TestGetEnumMap.ZapEnum}

tests/util.py

@@ -2,11 +2,15 @@
 import json
 import os
 import re
+from collections.abc import Callable
 from datetime import datetime, timezone
-from inspect import getcallargs, getfullargspec
+from enum import Enum
+from inspect import getcallargs, getfullargspec, signature
 from io import StringIO
 from os import path
 from pathlib import Path
+from types import UnionType
+from typing import Any, Dict, List, Set, Tuple, Union, get_args, get_type_hints, get_origin  # noqa: UP035
 from unittest.mock import patch
 
 from pyinfra.api import Config, Inventory
@@ -59,7 +63,31 @@ def get_temp_filename(*args):
         return "_tempfile_"
 
 
-def parse_value(value):
+AGGREGATES = {Dict, List, Set, Tuple, Union, UnionType, dict, list, set, tuple}  # noqa: UP006
+
+
+def get_enum_map(op: Callable[..., Any]) -> dict[str, type[Enum]]:
+    """
+    Returns a map from type name to type for all enum types used in `ops` parameters.
+    """
+    result: dict[str, type[Enum]] = {}
+    type_hints = get_type_hints(op)
+    for param in signature(op).parameters:
+        if param not in type_hints:
+            continue
+        to_do = [type_hints[param]]
+        while len(to_do) > 0:
+            the_type = to_do.pop(0)
+            origin = get_origin(the_type)
+            if (origin is not None) and (origin in AGGREGATES):
+                to_do.extend(get_args(the_type))
+            elif isinstance(the_type, type) and issubclass(the_type, Enum):
+                result[the_type.__name__] = the_type
+
+    return result
+
+
+def parse_value(value, enum_map: dict[str, type[Enum]] | None = None):
     """
     Convert JSON types to more complex Python types because JSON is lacking.
     """
@@ -69,17 +97,29 @@ def parse_value(value):
             return datetime.fromisoformat(value[9:])
         if value.startswith("path:"):
             return Path(value[5:])
+        if value.startswith("enum:"):
+            if len(pieces := value.split(":")) != 3:  # enum:<enum_type_name>:<value>
+                raise ValueError(f"invalid enum specifier: {value}")
+            try:
+                result = (enum_map or {})[pieces[1]](pieces[2])
+            except KeyError:
+                raise ValueError(f"enum '{pieces[1]}' not defined for '{value}'") from None
+            except ValueError:
+                raise ValueError(f"value '{pieces[2]}' not valid in '{value}'") from None
+            else:
+                return result
+
         if value.startswith("io:"):
             return StringIO(value[3:])
         return value
 
     if isinstance(value, list):
         if value and value[0] == "set:":
-            return set(parse_value(value) for value in value[1:])
-        return [parse_value(value) for value in value]
+            return set(parse_value(value, enum_map) for value in value[1:])
+        return [parse_value(value, enum_map) for value in value]
 
     if isinstance(value, dict):
-        return {key: parse_value(value) for key, value in value.items()}
+        return {key: parse_value(value, enum_map) for key, value in value.items()}
 
     return value