diff options
-rw-r--r-- | libmproxy/flow.py | 10 | ||||
-rw-r--r-- | test/test_flow.py | 11 |
2 files changed, 12 insertions, 9 deletions
diff --git a/libmproxy/flow.py b/libmproxy/flow.py index e0013f1e..4c778c40 100644 --- a/libmproxy/flow.py +++ b/libmproxy/flow.py @@ -1394,9 +1394,10 @@ class FlowMaster(controller.Master): """ pass - def unload_script(self, script): - script.unload() - self.scripts.remove(script) + def unload_scripts(self): + for script in self.scripts[:]: + script.unload() + self.scripts.remove(script) def load_script(self, command): """ @@ -1627,8 +1628,7 @@ class FlowMaster(controller.Master): return f def shutdown(self): - for script in self.scripts: - self.unload_script(script) + self.unload_scripts() controller.Master.shutdown(self) if self.stream: for i in self.state._flow_list: diff --git a/test/test_flow.py b/test/test_flow.py index 680d59e5..f9198f0c 100644 --- a/test/test_flow.py +++ b/test/test_flow.py @@ -565,8 +565,7 @@ class TestFlowMaster: fm = flow.FlowMaster(None, s) assert not fm.load_script(tutils.test_data.path("scripts/a.py")) assert not fm.load_script(tutils.test_data.path("scripts/a.py")) - assert not fm.unload_script(fm.scripts[0]) - assert not fm.unload_script(fm.scripts[0]) + assert not fm.unload_scripts() assert fm.load_script("nonexistent") assert "ValueError" in fm.load_script(tutils.test_data.path("scripts/starterr.py")) assert len(fm.scripts) == 0 @@ -613,9 +612,13 @@ class TestFlowMaster: fm.handle_clientdisconnect(dc) assert fm.scripts[0].ns["log"][-1] == "clientdisconnect" assert fm.scripts[1].ns["log"][-1] == "clientdisconnect" + + #unload first script - fm.unload_script(fm.scripts[0]) - assert len(fm.scripts) == 1 + fm.unload_scripts() + assert len(fm.scripts) == 0 + + assert not fm.load_script(tutils.test_data.path("scripts/all.py")) err = flow.Error(f.request, "msg") err.reply = controller.DummyReply() fm.handle_error(err) |