diff --git a/test/infamy/__init__.py b/test/infamy/__init__.py index fae02eb0..4c44ec8c 100644 --- a/test/infamy/__init__.py +++ b/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 diff --git a/test/infamy/env.py b/test/infamy/env.py index 8fd6032f..3cc85465 100644 --- a/test/infamy/env.py +++ b/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: