Fixed use of decorator with arguments

RCUSCC/RCUSCC/wrappers.py

@@ -4,20 +4,23 @@ import traceback
 
 __all__ = ["only_in_states", "only_when_on", "fault_on_error"]
 
-def only_in_states(func, allowed_states):
+def only_in_states(allowed_states):
     """
       Wrapper to call and return the wrapped function if the device is
       in one of the given states. Otherwise None is returned and nothing
       will be called.
     """
-    @wraps(func)
-    def state_check_wrapper(self, *args, **kwargs):
-        if self.get_state() in allowed_states:
-            return func(self, *args, **kwargs)
+    def wrapper(func):
+        @wraps(func)
+        def state_check_wrapper(self, *args, **kwargs):
+            if self.get_state() in allowed_states:
+                return func(self, *args, **kwargs)
 
-        return None
+            return None
+
+        return state_check_wrapper
 
-    return state_check_wrapper
+    return wrapper
 
 def only_when_on(func):
     """