Flatten ServerMaster into Master

This commit is contained in:
Aldo Cortesi 2016-05-29 12:01:46 +12:00
parent ed415877d4
commit e4f043f69c
3 changed files with 11 additions and 38 deletions

View File

@ -12,17 +12,18 @@ class ControlError(Exception):
class Master(object): class Master(object):
""" """
The master handles mitmproxy's main event loop. The master handles mitmproxy's main event loop.
""" """
def __init__(self): def __init__(self):
self.event_queue = queue.Queue() self.event_queue = queue.Queue()
self.should_exit = threading.Event() self.should_exit = threading.Event()
self.servers = []
def start(self): def start(self):
self.should_exit.clear() self.should_exit.clear()
for server in self.servers:
ServerThread(server).start()
def run(self): def run(self):
self.start() self.start()
@ -57,38 +58,18 @@ class Master(object):
return changed return changed
def shutdown(self): def shutdown(self):
for server in self.servers:
server.shutdown()
self.should_exit.set() self.should_exit.set()
class ServerMaster(Master):
"""
The ServerMaster adds server thread support to the master.
"""
def __init__(self):
super(ServerMaster, self).__init__()
self.servers = []
def add_server(self, server): def add_server(self, server):
# We give a Channel to the server which can be used to communicate with the master # We give a Channel to the server which can be used to communicate with the master
channel = Channel(self.event_queue, self.should_exit) channel = Channel(self.event_queue, self.should_exit)
server.set_channel(channel) server.set_channel(channel)
self.servers.append(server) self.servers.append(server)
def start(self):
super(ServerMaster, self).start()
for server in self.servers:
ServerThread(server).start()
def shutdown(self):
for server in self.servers:
server.shutdown()
super(ServerMaster, self).shutdown()
class ServerThread(threading.Thread): class ServerThread(threading.Thread):
def __init__(self, server): def __init__(self, server):
self.server = server self.server = server
super(ServerThread, self).__init__() super(ServerThread, self).__init__()
@ -100,12 +81,10 @@ class ServerThread(threading.Thread):
class Channel(object): class Channel(object):
""" """
The only way for the proxy server to communicate with the master The only way for the proxy server to communicate with the master
is to use the channel it has been given. is to use the channel it has been given.
""" """
def __init__(self, q, should_exit): def __init__(self, q, should_exit):
self.q = q self.q = q
self.should_exit = should_exit self.should_exit = should_exit
@ -141,12 +120,10 @@ class Channel(object):
class DummyReply(object): class DummyReply(object):
""" """
A reply object that does nothing. Useful when we need an object to seem A reply object that does nothing. Useful when we need an object to seem
like it has a channel, and during testing. like it has a channel, and during testing.
""" """
def __init__(self): def __init__(self):
self.acked = False self.acked = False
self.taken = False self.taken = False
@ -192,7 +169,6 @@ def handler(f):
class Reply(object): class Reply(object):
""" """
Messages sent through a channel are decorated with a "reply" attribute. Messages sent through a channel are decorated with a "reply" attribute.
This object is used to respond to the message through the return This object is used to respond to the message through the return

View File

@ -638,7 +638,7 @@ class State(object):
self.flows.kill_all(master) self.flows.kill_all(master)
class FlowMaster(controller.ServerMaster): class FlowMaster(controller.Master):
@property @property
def server(self): def server(self):

View File

@ -16,7 +16,6 @@ class TMsg:
class TestMaster(object): class TestMaster(object):
def test_simple(self): def test_simple(self):
class DummyMaster(controller.Master): class DummyMaster(controller.Master):
@controller.handler @controller.handler
def handle_panic(self, _): def handle_panic(self, _):
@ -34,10 +33,8 @@ class TestMaster(object):
m.run() m.run()
assert m.should_exit.is_set() assert m.should_exit.is_set()
def test_server_simple(self):
class TestServerMaster(object): m = controller.Master()
def test_simple(self):
m = controller.ServerMaster()
s = DummyServer(None) s = DummyServer(None)
m.add_server(s) m.add_server(s)
m.start() m.start()