Wrap a few more functions from salt.modules.state

Signed-off-by: Pedro Algarvio <palgarvio@vmware.com>
saltstack · Nov 4, 2022 · 83144c3 · 83144c3
1 parent 708bbde
commit 83144c3
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 2 deletions.
diff --git a/changelog/140.improvement.rst b/changelog/140.improvement.rst
@@ -1 +1,4 @@
-Allow getting the state chunk by `__id__` on MultiStateResult
+Several improvements to the state module wrappers:
+
+* Allow getting the state chunk by `__id__` on MultiStateResult
+* Wrap a few more functions from `salt.modules.state`
diff --git a/src/saltfactories/utils/functional.py b/src/saltfactories/utils/functional.py
@@ -181,12 +181,19 @@ def __setitem__(self, key, value):
                     wrapper which will return a more pythonic data structure to assert against.
                     """
                     if key in (
+                        "state.apply",
+                        "state.high",
+                        "state.highstate",
+                        "state.low",
                         "state.single",
                         "state.sls",
+                        "state.sls_id",
                         "state.template",
                         "state.template_str",
+                        "state.test",
+                        "state.top",
                     ):
-                        if key == "state.single":
+                        if key in ("state.low", "state.single", "state.sls_id"):
                             wrapper_cls = StateResult
                         else:
                             wrapper_cls = MultiStateResult