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,
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 __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 test_dsp_agent_connection(dispatcher):
with RequestConnection(port=DISPATCHER_PORT) as request_connection:
request_connection.establish()
logger.info('Sending test message')
dispatcher.request_handler = dummy_request_handler
reply = request_connection.send(TEST_MESSAGE, 1)
assert reply == TEST_MESSAGE, 'Test message in request is modified ' \
'between "request_handler" and "reply"'
def test_dsp_pulse(dispatcher):
with RequestConnection(port=DISPATCHER_PORT) as request_connection:
request_connection.establish()
register_req = deepcopy(Register_agent)
register_req['token'] = CLIENT_TEST_TOKEN
reply = request_connection.send(register_req, 1)
pulse_req = deepcopy(Pulse)
pulse_req['id'] = reply['id']
for _ in range(10):
reply = request_connection.send(pulse_req, 1)
assert reply['result'], 'Wrong reply status'
def test_dsp_client_queue(dispatcher):
with RequestConnection(port=DISPATCHER_PORT) as request_connection:
request_connection.establish()
request = deepcopy(Client_queues)
request['name'] = 'test_dsp_client_queue'
request['token'] = CLIENT_TEST_TOKEN
reply = request_connection.send(request, 1)
assert reply['name'] == request['name'], \
'Name param is modified or wrong reply'
# TODO: Validate broker host
assert reply['broker']['result'][QUEUE] == request['name'], \
'Wrong result queue name is defined by dispatcher'
assert reply['broker']['task'][QUEUE] == RoutingKeys.TASK, \
f'Task queue is not "{RoutingKeys.TASK}"'
def test_dsp_register(dispatcher):
name = 'this_is_test'
with RequestConnection(port=DISPATCHER_PORT) as request_connection:
request_connection.establish()
logger.info('Sending registration message')
register_req = deepcopy(Register_agent)
register_req['name'] = name
register_req['token'] = CLIENT_TEST_TOKEN
expected_id = dispatcher._next_free_id
reply = request_connection.send(register_req, 1)
assert reply['result'], 'Registration was not successful'
assert reply['id'] == expected_id, 'Wrong agent id is assigned'
assert expected_id in dispatcher.agents, 'Agent is missing in ' \
'dispatcher agents list'
assert dispatcher.agents[expected_id].name == name, 'Agent name does not ' \
'corresponds to expected one'
now = datetime.utcnow()
last_sync = dispatcher.agents[expected_id].last_sync
assert 0.01 > (now - last_sync).seconds, 'Request-Reply sync timestamp' \
'differs more than expected'
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)