From b7c1d057828054a7e8d5f9ccf2c3c93ca8888ed3 Mon Sep 17 00:00:00 2001
From: Vyacheslav Bakhmutov <m0sth8@yandex-team.ru>
Date: Fri, 13 Jun 2014 14:14:55 +0700
Subject: [PATCH] Remove global should_exit and fix tests

---
 libmproxy/console/__init__.py |  4 ++--
 libmproxy/controller.py       | 21 +++++++++------------
 libmproxy/flow.py             |  5 +++--
 test/test_flow.py             |  3 +--
 test/tservers.py              |  2 +-
 5 files changed, 16 insertions(+), 19 deletions(-)

diff --git a/libmproxy/console/__init__.py b/libmproxy/console/__init__.py
index c1ef43319..e660f312b 100644
--- a/libmproxy/console/__init__.py
+++ b/libmproxy/console/__init__.py
@@ -577,7 +577,7 @@ class ConsoleMaster(flow.FlowMaster):
 
         self.view_flowlist()
 
-        self.server.start_slave(controller.Slave, controller.Channel(self.masterq))
+        self.server.start_slave(controller.Slave, controller.Channel(self.masterq, self.should_exit))
 
         if self.options.rfile:
             ret = self.load_flows(self.options.rfile)
@@ -780,7 +780,7 @@ class ConsoleMaster(flow.FlowMaster):
     def loop(self):
         changed = True
         try:
-            while not controller.should_exit:
+            while not self.should_exit.is_set():
                 startloop = time.time()
                 if changed:
                     self.statusbar.redraw()
diff --git a/libmproxy/controller.py b/libmproxy/controller.py
index c4a608545..63e442411 100644
--- a/libmproxy/controller.py
+++ b/libmproxy/controller.py
@@ -1,9 +1,6 @@
 from __future__ import absolute_import
 import Queue, threading
 
-should_exit = False
-
-
 class DummyReply:
     """
         A reply object that does nothing. Useful when we need an object to seem
@@ -37,8 +34,9 @@ class Reply:
 
 
 class Channel:
-    def __init__(self, q):
+    def __init__(self, q, should_exit):
         self.q = q
+        self.should_exit = should_exit
 
     def ask(self, mtype, m):
         """
@@ -47,7 +45,7 @@ class Channel:
         """
         m.reply = Reply(m)
         self.q.put((mtype, m))
-        while not should_exit:
+        while not self.should_exit.is_set():
             try:
                 # The timeout is here so we can handle a should_exit event.
                 g = m.reply.q.get(timeout=0.5)
@@ -89,6 +87,7 @@ class Master:
         """
         self.server = server
         self.masterq = Queue.Queue()
+        self.should_exit = threading.Event()
 
     def tick(self, q):
         changed = False
@@ -107,10 +106,9 @@ class Master:
         return changed
 
     def run(self):
-        global should_exit
-        should_exit = False
-        self.server.start_slave(Slave, Channel(self.masterq))
-        while not should_exit:
+        self.should_exit.clear()
+        self.server.start_slave(Slave, Channel(self.masterq, self.should_exit))
+        while not self.should_exit.is_set():
             self.tick(self.masterq)
         self.shutdown()
 
@@ -123,8 +121,7 @@ class Master:
             obj.reply()
 
     def shutdown(self):
-        global should_exit
-        if not should_exit:
-            should_exit = True
+        if not self.should_exit.is_set():
+            self.should_exit.set()
             if self.server:
                 self.server.shutdown()
diff --git a/libmproxy/flow.py b/libmproxy/flow.py
index a440b850b..b6b490222 100644
--- a/libmproxy/flow.py
+++ b/libmproxy/flow.py
@@ -654,6 +654,7 @@ class FlowMaster(controller.Master):
                     self.server.config,
                     f,
                     self.masterq,
+                    self.should_exit
                 )
             rt.start() # pragma: no cover
             if block:
@@ -792,8 +793,8 @@ class FilteredFlowWriter:
 class RequestReplayThread(threading.Thread):
     name="RequestReplayThread"
 
-    def __init__(self, config, flow, masterq):
-        self.config, self.flow, self.channel = config, flow, controller.Channel(masterq)
+    def __init__(self, config, flow, masterq, should_exit):
+        self.config, self.flow, self.channel = config, flow, controller.Channel(masterq, should_exit)
         threading.Thread.__init__(self)
 
     def run(self):
diff --git a/test/test_flow.py b/test/test_flow.py
index 2365c08c9..890130035 100644
--- a/test/test_flow.py
+++ b/test/test_flow.py
@@ -672,7 +672,6 @@ class TestFlowMaster:
         fm.handle_error(f.error)
 
     def test_server_playback(self):
-        controller.should_exit = False
         s = flow.State()
 
         f = tutils.tflow()
@@ -695,7 +694,7 @@ class TestFlowMaster:
         fm.start_server_playback(pb, False, [], True, False)
         q = Queue.Queue()
         fm.tick(q)
-        assert controller.should_exit
+        assert fm.should_exit.is_set()
 
         fm.stop_server_playback()
         assert not fm.server_playback
diff --git a/test/tservers.py b/test/tservers.py
index 7548ebc4d..146dd13de 100644
--- a/test/tservers.py
+++ b/test/tservers.py
@@ -283,7 +283,7 @@ class ChainProxTest(ProxTestBase):
     def teardownAll(cls):
         super(ChainProxTest, cls).teardownAll()
         for p in cls.chain:
-            p.tmaster.server.shutdown()
+            p.tmaster.shutdown()
 
     def setUp(self):
         super(ChainProxTest, self).setUp()