diff options
Diffstat (limited to 'mitmproxy/test/taddons.py')
-rw-r--r-- | mitmproxy/test/taddons.py | 18 |
1 files changed, 15 insertions, 3 deletions
diff --git a/mitmproxy/test/taddons.py b/mitmproxy/test/taddons.py index 3cba6762..7804b90d 100644 --- a/mitmproxy/test/taddons.py +++ b/mitmproxy/test/taddons.py @@ -1,7 +1,10 @@ +import contextlib + import mitmproxy.master import mitmproxy.options from mitmproxy import proxy from mitmproxy import events +from mitmproxy import exceptions class RecordingMaster(mitmproxy.master.Master): @@ -36,6 +39,15 @@ class context: self.wrapped = None return False + @contextlib.contextmanager + def _rollback(self, opts, updates): + old = opts._opts.copy() + try: + yield + except exceptions.OptionsError as e: + opts.__dict__["_opts"] = old + raise + def cycle(self, addon, f): """ Cycles the flow through the events for the flow. Stops if a reply @@ -55,6 +67,6 @@ class context: Options object with the given keyword arguments, then calls the configure method on the addon with the updated value. """ - for k, v in kwargs.items(): - setattr(self.options, k, v) - addon.configure(self.options, kwargs.keys()) + with self._rollback(self.options, kwargs): + self.options.update(**kwargs) + addon.configure(self.options, kwargs.keys()) |