test: infamy: use decorator to reduce boilerplate for test args

Example use:

    @infamy.test_argument("--mode", help="one of 'static' or 'lacp'")
    class TestArgs(infamy.ArgumentParser):
        pass

Signed-off-by: Joachim Wiberg <troglobit@gmail.com>

test/infamy/__init__.py

@@ -3,6 +3,7 @@ import os
 from .container import Container
 from .env import Env
 from .env import ArgumentParser
+from .env import test_argument
 from .furl import Furl
 from .netns import IsolatedMacVlan,IsolatedMacVlans
 from .sniffer import Sniffer

test/infamy/env.py

@@ -16,6 +16,7 @@ class NullEnv:
 
 ENV = NullEnv()
 
+
 class ArgumentParser():
     def DefaultTransport():
         """Pick pseudo-random transport
@@ -41,7 +42,6 @@ class ArgumentParser():
         self.args.add_argument("ptop", nargs=1, metavar="topology")
         self.args.add_argument("-l", "--logical-topology", dest="ltop", default=top)
 
-
     def add_argument(self, *args, **kwargs):
         kwargs["required"] = True
         self.args.add_argument(*args, **kwargs)
@@ -49,6 +49,21 @@ class ArgumentParser():
     def parse_args(self, argv):
         return self.args.parse_args(argv)
 
+
+def test_argument(option, **kwargs):
+    """See lag_failure/test.py for an example @infamy.test_argumet()"""
+    def decorator(cls):
+        super_init = cls.__init__
+
+        def new_init(self, *args, **kw):
+            super_init(self, *args, **kw)
+            self.add_argument(option, **kwargs)
+
+        cls.__init__ = new_init
+        return cls
+    return decorator
+
+
 class Env(object):
     def __init__(self, ltop=None, args=None, argv=sys.argv[1::], environ=os.environ):
         if "INFAMY_ARGS" in environ: