From fdf6a6570c4342807525bf3b7401d491034e31f6 Mon Sep 17 00:00:00 2001
From: Darko Poljak <darko.poljak@gmail.com>
Date: Mon, 15 Aug 2016 16:01:39 +0200
Subject: [PATCH] Check for beta in scripts/cdist.

---
 cdist/__init__.py | 14 ++++++++++++++
 cdist/config.py   | 19 +------------------
 scripts/cdist     | 21 +++++++++++++++++++++
 3 files changed, 36 insertions(+), 18 deletions(-)

diff --git a/cdist/__init__.py b/cdist/__init__.py
index 74db1a13..d06ae28b 100644
--- a/cdist/__init__.py
+++ b/cdist/__init__.py
@@ -56,6 +56,20 @@ class UnresolvableRequirementsError(cdist.Error):
     pass
 
 
+class CdistBetaRequired(cdist.Error):
+    """Beta functionality is used but beta is not enabled"""
+
+    def __init__(self, command, arg):
+        self.command = command
+        self.arg = arg
+
+    def __str__(self):
+        err_msg = ("\'{}\' argument of \'{}\' command is beta, but beta is "
+                   "not enabled. If you want to use it please enable beta "
+                   "functionalities.")
+        return err_msg.format(self.arg, self.command)
+
+
 class CdistObjectError(Error):
     """Something went wrong with an object"""
 
diff --git a/cdist/config.py b/cdist/config.py
index ad256b12..9d4c5c10 100644
--- a/cdist/config.py
+++ b/cdist/config.py
@@ -71,8 +71,6 @@ def inspect_ssh_mux_opts():
 class Config(object):
     """Cdist main class to hold arbitrary data"""
 
-    BETA_ARGS = ['jobs']
-
     def __init__(self, local, remote, dry_run=False, jobs=None):
 
         self.local = local
@@ -111,19 +109,6 @@ class Config(object):
                 for host in source:
                     yield host
 
-    @classmethod
-    def _check_beta(cls, args_dict):
-        if 'beta' not in args_dict:
-            args_dict['beta'] = False
-        # Check only if beta is not enabled: if beta option is specified then
-        # raise error.
-        if not args_dict['beta']:
-            err_msg = ("\'{}\' is beta, but beta is not enabled. If you want "
-                       "to use it please enable beta functionalities.")
-            for arg in cls.BETA_ARGS:
-                if arg in args_dict:
-                    raise cdist.Error(err_msg.format(arg))
-
     @classmethod
     def commandline(cls, args):
         """Configure remote system"""
@@ -136,9 +121,6 @@ class Config(object):
             raise cdist.Error(("Cannot read both, manifest and host file, "
                                "from stdin"))
 
-        args_dict = vars(args)
-        cls._check_beta(args_dict)
-
         # if no host source is specified then read hosts from stdin
         if not (args.hostfile or args.host):
             args.hostfile = '-'
@@ -167,6 +149,7 @@ class Config(object):
         args.remote_exec_pattern = None
         args.remote_copy_pattern = None
 
+        args_dict = vars(args)
         # if remote-exec and/or remote-copy args are None then user
         # didn't specify command line options nor env vars:
         # inspect multiplexing options for default cdist.REMOTE_COPY/EXEC
diff --git a/scripts/cdist b/scripts/cdist
index 6c60ccb3..a2f5fbfb 100755
--- a/scripts/cdist
+++ b/scripts/cdist
@@ -22,6 +22,12 @@
 #
 
 
+# list of beta arguments for sub-commands
+BETA_ARGS = {
+    'config': ['jobs', ],
+}
+
+
 def check_positive_int(value):
     import argparse
 
@@ -36,6 +42,20 @@ def check_positive_int(value):
     return val
 
 
+def check_beta(args_dict):
+    if 'beta' not in args_dict:
+        args_dict['beta'] = False
+    # Check only if beta is not enabled: if beta option is specified then
+    # raise error.
+    if not args_dict['beta']:
+        err_msg = ("\'{}\' is beta, but beta is not enabled. If you want "
+                   "to use it please enable beta functionalities.")
+        cmd = args_dict['command']
+        for arg in BETA_ARGS[cmd]:
+            if arg in args_dict:
+                raise cdist.CdistBetaRequired(cmd, arg)
+
+
 def commandline():
     """Parse command line"""
     import argparse
@@ -172,6 +192,7 @@ def commandline():
         parser['main'].print_help()
         sys.exit(0)
 
+    check_beta(vars(args))
     args.func(args)
 
 if __name__ == "__main__":