testsuites/vstf/vstf_scripts/vstf/common/decorator.py

   1 ##############################################################################
   2 # Copyright (c) 2015 Huawei Technologies Co.,Ltd and others.
   3 #
   4 # All rights reserved. This program and the accompanying materials
   5 # are made available under the terms of the Apache License, Version 2.0
   6 # which accompanies this distribution, and is available at
   7 # http://www.apache.org/licenses/LICENSE-2.0
   8 ##############################################################################
   9
  10 _DEFAULTS = "vstf check key defaults".encode()
  11
  12
  13 def check(key, choices=[], defaults=_DEFAULTS):
  14     def _deco(func):
  15         def __deco(*args, **kwargs):
  16             if key not in kwargs:
  17                 if defaults != _DEFAULTS:
  18                     kwargs[key] = defaults
  19                 else:
  20                     raise Exception(
  21                         "Error: '%s' is needed in %s" %
  22                         (key, func))
  23
  24             if choices and kwargs[key] not in choices:
  25                 raise Exception("Error: %s :%s" % (key, kwargs[key]))
  26             ret = func(*args, **kwargs)
  27             return ret
  28
  29         return __deco
  30
  31     return _deco
  32
  33
  34 def dcheck(key, choices=[]):
  35     def _deco(func):
  36         def __deco(*args):
  37             if len(args) == 2:
  38                 values = args[1]
  39             elif len(args) == 1:
  40                 values = args[0]
  41             else:
  42                 values = None
  43             if isinstance(values, dict):
  44                 if key not in values:
  45                     raise Exception(
  46                         "Error: '%s' is needed in %s" %
  47                         (key, func))
  48                 if choices and values[key] not in choices:
  49                     raise Exception("Error: %s :%s" % (key, values[key]))
  50             ret = func(*args)
  51             return ret
  52
  53         return __deco
  54
  55     return _deco
  56
  57
  58 def vstf_input(key, types=str, choices=[], default=None):
  59     def _deco(func):
  60         def __deco(*args):
  61             ret = func(*args)
  62             if not isinstance(ret, dict):
  63                 ret = {}
  64             in_str = "----> %s : " % key
  65             if choices:
  66                 in_str = "---- %s\n" % (str(choices)) + in_str
  67             while True:
  68                 if types == list or types == dict:
  69                     value = input(in_str)
  70                 else:
  71                     value = raw_input(in_str)
  72                     value = types(value)
  73                 if not value:
  74                     value = default
  75                 if not choices or value in choices:
  76                     break
  77             ret.update({key: value})
  78             return ret
  79
  80         return __deco
  81
  82     return _deco
  83
  84
  85 def namespace():
  86     def _deco(func):
  87         def __deco(*args, **kwargs):
  88             ret = func(*args, **kwargs)
  89             nspace = kwargs.get("namespace", None)
  90             if nspace:
  91                 ret = "ip netns exec %(namespace)s " % {
  92                     "namespace": nspace} + ret
  93             return ret
  94
  95         return __deco
  96
  97     return _deco