Handle test mode when enabling ufw

_modules/ufw.py

@@ -22,14 +22,19 @@ def get_default_outgoing():
     return policy
 
 def set_enabled(enabled):
-    cmd = 'ufw --force enable' if enabled else 'ufw disable'
+    if __opts__['test']:
+        cmd = "ufw --dry-run "
+    else:
+        cmd = "ufw "
+    cmd += '--force enable' if enabled else 'disable'
     __salt__['cmd.run'](cmd)
 
 def add_rule(rule):
     if __opts__['test']:
-        cmd = "ufw --dry-run " + rule
+        cmd = "ufw --dry-run "
     else:
-        cmd = "ufw " + rule
+        cmd = "ufw "
+    cmd += rule
     out = __salt__['cmd.run'](cmd, python_shell=True)
     return out
 

_states/ufw.py

@@ -84,15 +84,15 @@ def enabled(name, **kwargs):
     if __salt__['ufw.is_enabled']():
         return _unchanged(name, "UFW is already enabled")
 
-    if __opts__['test']:
-        return _test(name, "UFW will be enabled")
-
     try:
         __salt__['ufw.set_enabled'](True)
     except (CommandExecutionError, CommandNotFoundError) as e:
         return _error(name, e.message)
 
-    return _changed(name, "UFW is enabled", enabled=True)
+    if __opts__['test']:
+        return _test(name, "UFW would have been enabled")
+    else:
+        return _changed(name, "UFW is enabled", enabled=True)
 
 
 def default_incoming(name, default):