Fixed use of decorator with arguments

RCUSCC/RCUSCC/wrappers.py

@@ -4,12 +4,13 @@ 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.
     """
+    def wrapper(func):
         @wraps(func)
         def state_check_wrapper(self, *args, **kwargs):
             if self.get_state() in allowed_states:
@@ -19,6 +20,8 @@ def only_in_states(func, allowed_states):
 
         return state_check_wrapper
 
+    return wrapper
+
 def only_when_on(func):
     """
       Wrapper to call and return the wrapped function if the device is