def _many_message(self, count, drop_count, command): if type(command) == list: in_command = os.linesep.join(command) count = len(command) else: in_command = "" for i in range(count): in_command = in_command + command + os.linesep inlines = io.StringIO(in_command) outfile = io.StringIO() conn = self._get_conn(inlines, outfile, drop_count) request_listener = reply.RequestListener(self.conf_obj, conn, self.disp, self.db) conn.set_receiver(request_listener) self.disp.start_workers(request_listener) # wait until the request is done while request_listener.get_messages_processed() < count: dcm_events.poll() for line in outfile.getvalue().split('\n'): line = line.strip() if line: output = json.loads(line.strip()) self.assertEquals(0, output['return_code']) request_listener.shutdown() request_listener.wait_for_all_nicely()
def test_reply_ack_timeout(self): conn = mock.Mock() reply_listener = mock.Mock() request_id = "requestID" message_id = "messageID" reply_payload = {"reply": "payload"} incoming_message = {"incoming": "info"} reply_rpc = reply.ReplyRPC(reply_listener, "AGENTID", conn, request_id, incoming_message, self.db, timeout=1) reply_rpc.reply(reply_payload) reply_doc = { "type": types.MessageTypes.ACK, "request_id": request_id, "message_id": message_id, "payload": reply_payload } dcm_events.poll(timeblock=1.2) reply_rpc.incoming_message(reply_doc) reply_listener.message_done.assert_called_once_with(reply_rpc) self.assertEqual(conn.send.call_count, 2)
def run_with_connect_errors(self, backoff_seconds, max_backoff_seconds, run_time_seconds, conn_obj): class FakeHS(object): def get_send_document(self): ws.throw_error(Exception("just for tests")) return {} def incoming_document(self, incoming_doc): return handshake.HandshakeIncomingReply( handshake.HandshakeIncomingReply.REPLY_CODE_SUCCESS) m = mock.Mock() conn_obj.return_value = m server_url = "wss://notreal.com" ws = websocket.WebSocketConnection( server_url, backoff_amount=int(backoff_seconds * 1000), max_backoff=int(max_backoff_seconds * 1000)) ws.connect(fake_incoming_message, FakeHS()) nw = datetime.datetime.now() done_time = nw + datetime.timedelta(seconds=run_time_seconds) while done_time > nw: remaining = done_time - nw dcm_events.poll(timeblock=remaining.total_seconds()) nw = datetime.datetime.now() ws.close() return m
def test_twosends_two_acks(self): timeout = 0.1 alert_doc = self._make_fake_alert_message() conn = mock.Mock() alerter = alert_msg.AlertAckMsg(alert_doc, conn, timeout=timeout) alerter.send() dcm_events.poll(timeblock=timeout * 1.5) self.assertGreaterEqual(conn.send.call_count, 2)
def test_twosends_two_acks(self): timeout = 0.1 alert_doc = self._make_fake_alert_message() conn = mock.Mock() alerter = alert_msg.AlertAckMsg(alert_doc, conn, timeout=timeout) alerter.send() dcm_events.poll(timeblock=timeout*1.5) self.assertGreaterEqual(conn.send.call_count, 2)
def test_alert_retransmission(self): timeout = 0.1 alert_doc = self._make_fake_alert_message() conn = mock.Mock() alerter = alert_msg.AlertAckMsg(alert_doc, conn, timeout=timeout) alerter.send() dcm_events.poll(timeblock=timeout * 1.5) alerter.incoming_message() call = mock.call(alert_doc) self.assertGreaterEqual(conn.send.call_count, 2) self.assertEqual(conn.send.call_args_list[0], call)
def test_alert_retransmission(self): timeout = 0.1 alert_doc = self._make_fake_alert_message() conn = mock.Mock() alerter = alert_msg.AlertAckMsg(alert_doc, conn, timeout=timeout) alerter.send() dcm_events.poll(timeblock=timeout*1.5) alerter.incoming_message() call = mock.call(alert_doc) self.assertGreaterEqual(conn.send.call_count, 2) self.assertEqual(conn.send.call_args_list[0], call)
def test_request_no_ack_timeout(self): conn = mock.Mock() doc = {'amessage': 'foru'} requester = request.RequestRPC(doc, conn, "XYZ", timeout=1) requester.send() dcm_events.poll(timeblock=1.5) self.assertGreater(conn.send.call_count, 1) (param_list, keywords) = conn.send.call_args self._validate_request_message(param_list[0], doc) requester.cleanup()
def test_request_no_ack_timeout(self): conn = mock.Mock() doc = {'amessage': 'foru'} requester = request.RequestRPC(doc, conn, "XYZ", timeout=1) requester.send() dcm_events.poll(timeblock=1.5) self.assertGreater(conn.send.call_count, 1) (param_list, keywords) = conn.send.call_args self._validate_request_message(param_list[0], doc) requester.cleanup()
def test_standard_with_callback_path(self): self.called = False def reply_called(*args, **kwargs): self.called = True conn = mock.Mock() doc = {'amessage': 'foru'} requester = request.RequestRPC(doc, conn, "XYZ", reply_callback=reply_called) requester.send() self.assertEqual(conn.send.call_count, 1) (param_list, keywords) = conn.send.call_args send_doc = param_list[0] reply_doc = { 'type': types.MessageTypes.ACK, 'message_id': send_doc['message_id'], 'request_id': send_doc['request_id'] } requester.incoming_message(reply_doc) self.assertEqual('REQUESTED', requester._sm._current_state) reply_doc = { 'type': types.MessageTypes.REPLY, 'message_id': send_doc['message_id'], 'request_id': send_doc['request_id'] } requester.incoming_message(reply_doc) while requester._sm._current_state !=\ states.RequesterStates.ACK_SENT: dcm_events.poll() self.assertEqual(states.RequesterStates.ACK_SENT, requester._sm._current_state) (param_list, keywords) = conn.send.call_args send_doc = param_list[0] self.assertEqual(reply_doc['message_id'], send_doc['message_id']) self.assertTrue('request_id' in send_doc) self.assertTrue('type' in send_doc) self.assertEqual(send_doc['type'], types.MessageTypes.ACK) requester.ack_sent_timeout()
def test_force_backoff(self, conn_obj): # force the backoff to be longer than the max run time then make sure # that the connect is only called once backoff_seconds = 0.2 max_backoff_seconds = backoff_seconds run_time_seconds = backoff_seconds * 10.0 force_time = run_time_seconds + 1.0 m = mock.Mock() conn_obj.return_value = m server_url = "wss://notreal.com" ws = websocket.WebSocketConnection( server_url, backoff_amount=int(backoff_seconds * 1000), max_backoff=int(max_backoff_seconds * 1000)) def send_in_handshake(): ws.event_incoming_message({ handshake.HandshakeIncomingReply.REPLY_KEY_FORCE_BACKOFF: force_time, 'return_code': handshake.HandshakeIncomingReply.REPLY_CODE_FORCE_BACKOFF }) class FakeHS(object): def get_send_document(self): dcm_events.register_callback(send_in_handshake) return {} def incoming_document(self, incoming_doc): hs = handshake.HandshakeIncomingReply( handshake.HandshakeIncomingReply.REPLY_CODE_FORCE_BACKOFF, force_backoff=force_time) return hs ws.connect(fake_incoming_message, FakeHS()) nw = datetime.datetime.now() done_time = nw + datetime.timedelta(seconds=run_time_seconds) while done_time > nw: remaining = done_time - nw dcm_events.poll(timeblock=remaining.total_seconds()) nw = datetime.datetime.now() ws.close() self.assertEqual(1, m.connect.call_count)
def _many_message(self, drop_count, command, retrans_list): class TestStateObserver(reply.ReplyObserverInterface): def new_message(self, reply): pass def message_done(self, reply): self.state_change_list = reply._sm.get_event_list() def incoming_message(self, incoming_doc): pass request_listener = None disp = dispatcher.Dispatcher(self.conf_obj) try: in_command = os.linesep.join(command) count = len(command) inlines = io.StringIO(in_command) outfile = io.StringIO() conn = self._get_conn(inlines, outfile, drop_count, retrans_list) request_listener = reply.RequestListener(self.conf_obj, conn, disp, self.db) conn.set_receiver(request_listener) to = TestStateObserver() rol = request_listener.get_reply_observers() rol.insert(0, to) disp.start_workers(request_listener) # wait until the request is done. in the case of reply # retransmissions this value could be greater than count while request_listener.get_messages_processed() < count: dcm_events.poll() for line in outfile.getvalue().split('\n'): line = line.strip() if line: output = json.loads(line.strip()) self.assertEquals(0, output['return_code']) finally: if request_listener: request_listener.shutdown() request_listener.wait_for_all_nicely() if disp: disp.stop() return to.state_change_list
def test_file_run(self): stat_name = str(uuid.uuid4()) hold_count = 5 check_interval = 0.1 # we are now setup for the test arguments = { "statType": "cpu-idle", "statName": stat_name, "holdCount": hold_count, "checkInterval": check_interval } plugin = init_system_stat.load_plugin(self.conf_obj, str(uuid.uuid4()), {}, "init_system_stat", arguments) result = plugin.run() result = result.get_reply_doc() self.assertEqual(result['return_code'], 0) try: event_space.poll(timeblock=check_interval * (hold_count + 1)) arguments = {"statName": stat_name} plugin = get_system_stat.load_plugin(self.conf_obj, str(uuid.uuid4()), {}, "get_system_stat", arguments) result = plugin.run() result = result.get_reply_doc() self.assertEqual(result['return_code'], 0) self.assertEqual(result['reply_type'], 'cpu_idle_stat_array') ro = result['reply_object'] self.assertEqual(len(ro['status']), hold_count) arguments = {"statName": stat_name} plugin = delete_system_stat.load_plugin(self.conf_obj, str(uuid.uuid4()), {}, "delete_system_stat", arguments) result = plugin.run() result = result.get_reply_doc() self.assertEqual(result['return_code'], 0) finally: try: systemstats.stop_stats(stat_name) except Exception as ex: print(str(ex))
def test_logging_handler_with_conn(self): conn = mock.Mock() conf = mock.Mock() logger_name = str(uuid.uuid4()) my_logger = logging.getLogger(logger_name) handler = logger.dcmLogger() my_logger.addHandler(handler) logger.set_dcm_connection(conf, conn) msg = "Test message with conn" my_logger.error(msg) handler.flush() dcm_events.poll(timeblock=0.0) args, kwargs = conn.send.call_args log_dict = args[0] self.assertEqual(log_dict['type'], "LOG") self.assertEqual(log_dict['level'], "ERROR") self.assertEqual(urllib.parse.unquote(log_dict['message']), msg)
def test_logging_handler_with_conn(self): conn = mock.Mock() conf = mock.Mock() logger_name = str(uuid.uuid4()) my_logger = logging.getLogger(logger_name) handler = logger.dcmLogger() my_logger.addHandler(handler) logger.set_dcm_connection(conf, conn) msg = "Test message with conn" my_logger.error(msg) handler.flush() dcm_events.poll(timeblock=0.0) args, kwargs = conn.send.call_args log_dict = args[0] self.assertEqual(log_dict["type"], "LOG") self.assertEqual(log_dict["level"], "ERROR") self.assertEqual(urllib.parse.unquote(log_dict["message"]), msg)
def test_standard_with_callback_path(self): self.called = False def reply_called(*args, **kwargs): self.called = True conn = mock.Mock() doc = {'amessage': 'foru'} requester = request.RequestRPC( doc, conn, "XYZ", reply_callback=reply_called) requester.send() self.assertEqual(conn.send.call_count, 1) (param_list, keywords) = conn.send.call_args send_doc = param_list[0] reply_doc = {'type': types.MessageTypes.ACK, 'message_id': send_doc['message_id'], 'request_id': send_doc['request_id']} requester.incoming_message(reply_doc) self.assertEqual('REQUESTED', requester._sm._current_state) reply_doc = {'type': types.MessageTypes.REPLY, 'message_id': send_doc['message_id'], 'request_id': send_doc['request_id']} requester.incoming_message(reply_doc) while requester._sm._current_state !=\ states.RequesterStates.ACK_SENT: dcm_events.poll() self.assertEqual(states.RequesterStates.ACK_SENT, requester._sm._current_state) (param_list, keywords) = conn.send.call_args send_doc = param_list[0] self.assertEqual(reply_doc['message_id'], send_doc['message_id']) self.assertTrue('request_id' in send_doc) self.assertTrue('type' in send_doc) self.assertEqual(send_doc['type'], types.MessageTypes.ACK) requester.ack_sent_timeout()
def test_file_run(self): stat_name = str(uuid.uuid4()) hold_count = 5 check_interval = 0.1 # we are now setup for the test arguments = {"statType": "cpu-idle", "statName": stat_name, "holdCount": hold_count, "checkInterval": check_interval} plugin = init_system_stat.load_plugin( self.conf_obj, str(uuid.uuid4()), {}, "init_system_stat", arguments) result = plugin.run() result = result.get_reply_doc() self.assertEqual(result['return_code'], 0) try: event_space.poll(timeblock=check_interval*(hold_count+1)) arguments = {"statName": stat_name} plugin = get_system_stat.load_plugin( self.conf_obj, str(uuid.uuid4()), {}, "get_system_stat", arguments) result = plugin.run() result = result.get_reply_doc() self.assertEqual(result['return_code'], 0) self.assertEqual(result['reply_type'], 'cpu_idle_stat_array') ro = result['reply_object'] self.assertEqual(len(ro['status']), hold_count) arguments = {"statName": stat_name} plugin = delete_system_stat.load_plugin( self.conf_obj, str(uuid.uuid4()), {}, "delete_system_stat", arguments) result = plugin.run() result = result.get_reply_doc() self.assertEqual(result['return_code'], 0) finally: try: systemstats.stop_stats(stat_name) except Exception as ex: print(str(ex))
def _rpc_wait_reply(self, doc): def reply_callback(): pass reqRPC = request.RequestRPC(doc, self.req_conn, self.agent_id, reply_callback=reply_callback) self.reply_conn.set_request_side(reqRPC) reqRPC.send() # wait for message completion: while not self._event.isSet(): dcm_events.poll() self._event.clear() reqRPC.cleanup() self.shutting_down = False return reqRPC
def _rpc_wait_reply(self, doc): def reply_callback(): pass reqRPC = request.RequestRPC(doc, self.req_conn, self.agent_id, reply_callback=reply_callback) self.reply_conn.set_request_side(reqRPC) reqRPC.send() # wait for message completion: while not self._event.isSet(): dcm_events.poll() self._event.clear() reqRPC.cleanup() self.shutting_down = False return reqRPC
def tests_sweeper(self): json_entry = {"12345": "6789"} j_size = len(json.dumps(json_entry)) per_page = 3 page_size = per_page * j_size page_monitor = pages.PageMonitor( page_size=page_size, life_span=2, sweep_time=1) page_monitor.start() try: json_list = [json_entry, json_entry] token = "pudding" pager = pages.JsonPage(page_size, json_list) page_monitor.new_pager(pager, token) dcm_events.poll(timeblock=3.0) self.assertRaises(exceptions.AgentPageNotFoundException, page_monitor.get_next_page, token) finally: page_monitor.stop()
def tests_sweeper(self): json_entry = {"12345": "6789"} j_size = len(json.dumps(json_entry)) per_page = 3 page_size = per_page * j_size page_monitor = pages.PageMonitor(page_size=page_size, life_span=2, sweep_time=1) page_monitor.start() try: json_list = [json_entry, json_entry] token = "pudding" pager = pages.JsonPage(page_size, json_list) page_monitor.new_pager(pager, token) dcm_events.poll(timeblock=3.0) self.assertRaises(exceptions.AgentPageNotFoundException, page_monitor.get_next_page, token) finally: page_monitor.stop()
def _simple_message(self, drop_count, command, stdout, stderr): inlines = io.StringIO(command) outfile = io.StringIO() conn = self._get_conn(inlines, outfile, drop_count) request_listener = reply.RequestListener(self.conf_obj, conn, self.disp, self.db) try: conn.set_receiver(request_listener) self.disp.start_workers(request_listener) # wait until the request is done while request_listener.get_messages_processed() != 1: dcm_events.poll() line = outfile.getvalue().split('\n')[0] line = line.strip() output = json.loads(line) self.assertEquals(stdout, output['message'].strip()) self.assertEquals(stderr, output['error_message']) self.assertEquals(0, output['return_code']) finally: request_listener.shutdown() request_listener.wait_for_all_nicely()
def test_reply_ack_timeout(self): conn = mock.Mock() reply_listener = mock.Mock() request_id = "requestID" message_id = "messageID" reply_payload = {"reply": "payload"} incoming_message = {"incoming": "info"} reply_rpc = reply.ReplyRPC( reply_listener, "AGENTID", conn, request_id, incoming_message, self.db, timeout=1) reply_rpc.reply(reply_payload) reply_doc = {"type": types.MessageTypes.ACK, "request_id": request_id, "message_id": message_id, "payload": reply_payload } dcm_events.poll(timeblock=1.2) reply_rpc.incoming_message(reply_doc) reply_listener.message_done.assert_called_once_with(reply_rpc) self.assertEqual(conn.send.call_count, 2)
def wait_for_all_nicely(self): # XXX TODO how long should this block? do we need this? # looks like just for tests while self._requests: dcm_events.poll()