def test_broker_smoke():
with Broker('localhost') as client:
client.connect()
client.setup_exchange()
client.declare(output_queue=task_input_queue)
logger.info('Sending task')
for i in test_tasks:
client.push(test_tasks[i])
with Broker('localhost') as agent:
agent.connect()
agent.declare(input_queue=task_input_queue,
output_queue=task_result_queue)
agent._inactivity_timeout = 0.1
logger.info('Processing task')
for task in agent.pulling_generator():
result = process_task(task)
agent.set_task_done(task)
agent.push(result)
with Broker('localhost') as validator:
validator.connect()
validator.declare(input_queue=task_result_queue)
validator._inactivity_timeout = 0.1
logger.info('Validating results')
for task in validator.pulling_generator():
validator_callback(task)
validator.set_task_done(task)
class Client:
def __init__(self,
name: str,
token: str,
dsp_ip: str = 'localhost',
dsp_port: int = 9999):
logger.info('Starting Client')
self.name = name
self.token = token
self.socket = RequestConnection(dsp_ip, dsp_port)
self.broker = None
def __enter__(self):
self.socket.establish()
return self
def __exit__(self, *exc_info):
self.socket.close()
def get_client_queues(self):
request = deepcopy(Client_queues)
request['name'] = self.name
request['token'] = self.token
reply = self.socket.send(request)
if reply['result']:
self.broker = Broker(reply['broker']['host'])
self.broker.connect()
self.broker.declare(reply['broker']['result'],
reply['broker']['task'])
else:
ConnectionRefusedError('Invalid credentials or resource is busy')
def __init__(self,
ip: str = '*',
port: Union[int, str] = '',
broker_host: str = ''):
logger.info('Starting Dispatcher')
self.socket = ReplyConnection(ip, port)
self.broker = Broker(broker_host if broker_host else ip)
self.agents = {}
self.request_handler = self.default_request_handler
self._next_free_id = INIT_AGENT_ID
self._listen = True
self._interrupt: Union[None, Callable] = None
def get_client_queues(self):
request = deepcopy(Client_queues)
request['name'] = self.name
request['token'] = self.token
reply = self.socket.send(request)
if reply['result']:
self.broker = Broker(reply['broker']['host'])
self.broker.connect()
self.broker.declare(reply['broker']['result'],
reply['broker']['task'])
else:
ConnectionRefusedError('Invalid credentials or resource is busy')
def init_broker(self):
"""
Request Agent queues on Broker from Dispatcher.
"""
request = deepcopy(Agent_queues)
request['token'] = self.token
request['id'] = self.id
reply = self.socket.send(request)
if reply['result']:
self.sync(reply)
self.broker = Broker(reply['broker']['host'])
self.broker.connect()
self.broker.declare(reply['broker']['task'],
reply['broker']['result'])
def broker():
host = '*'
with Broker(host) as broker:
broker.connect()
# broker._interrupt = polling_expiration
broker._inactivity_timeout = 0.1 * SECOND
broker.declare(input_queue=task_queue,
output_queue=compose_queue(RoutingKeys.RESULTS))
yield broker
input_queue = broker.input_queue
flush_queue(host, input_queue)
def flush_queue(broker: str,
queue: dict = task_queue,
assert_non_empty: bool = True):
with Broker(broker) as br:
br.connect()
br.declare(queue)
br._inactivity_timeout = 0.1 # seconds
empty = True
logger.info(f'Flushing queue: {queue}')
for task in br.pulling_generator():
empty = False
br.set_task_done(task)
if assert_non_empty:
assert empty, f'Flushed queue {queue} is not empty'
def test_agent_queues(agent_on_dispatcher: Agent, broker: Broker):
agent = agent_on_dispatcher
test_task = deepcopy(task_body)
test_task['arguments'] = 'agent_task_test'
task_result = deepcopy(task_body)
task_result['arguments'] = 'agent_task_result'
broker.declare(input_queue=compose_queue(RoutingKeys.RESULTS),
output_queue=compose_queue(RoutingKeys.TASK))
broker.push(test_task)
task = next(agent.broker.pulling_generator())
agent.broker.set_task_done(task)
assert test_task == task.body, \
'Wrong task is received from task queue for Agent'
agent.broker.push(task_result)
result = next(broker.pulling_generator())
broker.set_task_done(result)
assert task_result == result.body, \
'Wrong Agent result is received from task queue'
class Agent(AgentBase):
def __init__(self,
token: str,
dsp_host: str = 'localhost',
dsp_port: int = 9999):
logger.info('Starting Agent')
super(Agent, self).__init__()
self.socket = RequestConnection(dsp_host, dsp_port)
self.broker = None
self.token = token
def __enter__(self):
self.socket.establish()
return self
def __exit__(self, *exc_info):
self.close()
def close(self):
if self.broker:
self.broker.close()
self.socket.close()
def register(self):
request = deepcopy(Register_agent)
if self.name:
request['name'] = self.name
request['token'] = self.token
reply = self.socket.send(request)
if reply['result']:
self.id = reply['id']
self.sync(reply)
return reply['result']
def init_broker(self):
"""
Request Agent queues on Broker from Dispatcher.
"""
request = deepcopy(Agent_queues)
request['token'] = self.token
request['id'] = self.id
reply = self.socket.send(request)
if reply['result']:
self.sync(reply)
self.broker = Broker(reply['broker']['host'])
self.broker.connect()
self.broker.declare(reply['broker']['task'],
reply['broker']['result'])
def pulse(self) -> bool:
request = deepcopy(Pulse)
request['id'] = self.id
reply = self.socket.send(request)
self.sync(reply['reply'])
return reply['result']
def sync(self, reply: dict):
self.last_sync = datetime.utcnow()
if 'commands' in reply:
self.commands = reply['commands']
def apply_commands(self):
for command in self.commands:
_method = getattr(self, command)
if not _method():
logger.error(f'Method "{_method}" has failed to apply')
return False
else:
return True
def disconnect(self):
request = deepcopy(Disconnect)
request['id'] = self.id
reply = self.socket.send(request)
if reply['result']:
self.broker.close()
self.broker = None
self.id = 0
return reply['result']
def shutdown(self):
self.__exit__()
exit(0)
class Dispatcher:
def __init__(self,
ip: str = '*',
port: Union[int, str] = '',
broker_host: str = ''):
logger.info('Starting Dispatcher')
self.socket = ReplyConnection(ip, port)
self.broker = Broker(broker_host if broker_host else ip)
self.agents = {}
self.request_handler = self.default_request_handler
self._next_free_id = INIT_AGENT_ID
self._listen = True
self._interrupt: Union[None, Callable] = None
def __enter__(self):
self.socket.establish()
self.configure_broker()
return self
def __exit__(self, *exc_info):
logger.info(f'Closing Dispatcher connection:{self.socket}')
self.socket.close()
self.broker.close()
def configure_broker(self):
self.broker._inactivity_timeout = 0.01 * SECOND
if self.broker.connect() and not self.broker.input_queue:
self.broker.setup_exchange()
self.broker.declare(
input_queue=compose_queue(RoutingKeys.DISPATCHER))
def listen(self, polling_timeout: int = 10 * SECOND):
while self._listen:
expired = self.socket.listen(self.request_handler, polling_timeout)
if self._interrupt and self._interrupt(expired):
break
if not self.broker.input_queue and not self.broker.channel.is_open:
self.configure_broker()
else:
for task in self.broker.pulling_generator():
logger.debug(f'Got dispatcher task {task}')
def default_request_handler(self, request: dict):
commands = {
Commands.Register_agent: self._register_agent_handler,
Commands.Agent_queues: self._agent_queues_handler,
Commands.Pulse: self._pulse_handler,
Commands.Client_queues: self._client_handler,
Commands.Relay: self._relay,
Commands.Disconnect: self._disconnect_handler
}
assert request['command'] in commands, 'Command ' \
f'{request["command"]} is not registered in dispatcher request handler'
command = commands[request['command']]
return command(request)
def _register_agent_handler(self, request: dict):
logger.info(
f'Registration request received {request["name"]}({self._next_free_id})'
)
request['id'] = self._next_free_id
agent = RemoteAgent(self._next_free_id)
agent.name = request['name']
agent.token = request['token']
self.agents[self._next_free_id] = agent
logger.info(f'New agent id={agent.id}')
request['result'] = True
self._next_free_id += 1
return request
def _agent_queues_handler(self, request: dict):
"""
Returns Host and queues that agent should connect to.
"""
agent = self.agents[request['id']]
logger.info(f'Agent queues request received from {agent}')
config = Database.get_agent_param(agent.token)
request['broker']['host'] = config['broker']
request['broker']['task'] = compose_queue(RoutingKeys.TASK)
request['broker']['result'] = compose_queue(RoutingKeys.RESULTS)
request['result'] = True
return request
def _pulse_handler(self, request: dict):
logger.info(f'Pulse request received {request["id"]}')
agent = self.agents[request['id']]
reply = agent.sync(request)
return reply
def _client_handler(self, request: dict):
logger.info(f'Client queues are requested by: {request["name"]}')
config = Database.get_client_param(request['token'])
request['broker']['host'] = config['broker']
request['broker']['task'] = compose_queue(RoutingKeys.TASK)
request['broker']['result'] = compose_queue(request['name'])
request['result'] = True
return request
def _disconnect_handler(self, request: dict):
"""
Removes agent instance on dispatcher.
"""
logger.info(f'Disconnect request received {request["id"]}')
if request['id'] in self.agents:
self.agents.pop(request['id'])
request['result'] = True
return request
def _relay(self, request: dict):
"""
Returns received request.
"""
logger.debug(f'Relay request called')
return request