async def async_test_upload_to_host_inventory_via_kafka_send_exception(self): """Test uploading to inventory via kafka.""" self.processor.report_or_slice = self.report_slice hosts = {str(self.uuid): {'bios_uuid': str(self.uuid), 'name': 'value'}, str(self.uuid2): {'insights_client_id': 'value', 'name': 'foo'}, str(self.uuid3): {'ip_addresses': 'value', 'name': 'foo'}, str(self.uuid4): {'mac_addresses': 'value', 'name': 'foo'}, str(self.uuid5): {'vm_uuid': 'value', 'name': 'foo'}, str(self.uuid6): {'etc_machine_id': 'value'}, str(self.uuid7): {'subscription_manager_id': 'value'}} test_producer = AIOKafkaProducer( loop=report_slice_processor.SLICE_PROCESSING_LOOP, bootstrap_servers=report_slice_processor.INSIGHTS_KAFKA_ADDRESS ) # test KafkaConnectionException def raise_error(): """Raise a general error.""" raise Exception('Test') test_producer.start = CoroutineMock() test_producer.send = CoroutineMock(side_effect=raise_error) test_producer.stop = CoroutineMock() with self.assertRaises(msg_handler.KafkaMsgHandlerError): with patch('processor.report_slice_processor.AIOKafkaProducer', return_value=test_producer): await self.processor._upload_to_host_inventory_via_kafka(hosts)
async def async_test_upload_to_host_inventory_via_kafka(self): """Test uploading to inventory via kafka.""" self.processor.report_or_slice = self.report_slice hosts = { str(self.uuid): {'bios_uuid': str(self.uuid), 'name': 'value'}, str(self.uuid2): {'insights_client_id': 'value', 'name': 'foo'}, str(self.uuid3): {'ip_addresses': 'value', 'name': 'foo'}, str(self.uuid4): {'mac_addresses': 'value', 'name': 'foo'}, str(self.uuid5): {'vm_uuid': 'value', 'name': 'foo'}, str(self.uuid6): {'etc_machine_id': 'value'}, str(self.uuid7): {'subscription_manager_id': 'value'}, str(self.uuid8): {'system_profile': {'os_release': '7', 'os_kernel_version': '2.6.32'} }} test_producer = AIOKafkaProducer( loop=report_slice_processor.SLICE_PROCESSING_LOOP, bootstrap_servers=report_slice_processor.INSIGHTS_KAFKA_ADDRESS ) test_producer.start = CoroutineMock() test_producer.send = CoroutineMock() test_producer.stop = CoroutineMock() with patch('processor.report_slice_processor.AIOKafkaProducer', return_value=test_producer): with patch('processor.report_slice_processor.asyncio.wait', side_effect=None): # all though we are not asserting any results, the test here is # that no error was raised await self.processor._upload_to_host_inventory_via_kafka(hosts)
class KafkaProducer(BrokerInterface): def __init__(self, connection_string: str, **kwargs): self.producer = AIOKafkaProducer(bootstrap_servers=connection_string, **kwargs) self.started = False async def init_broker(self): await self.producer.start() self.started = True async def send(self, *, message: Dict, channel_name: str, timeout: int = 2, **kwargs): assert self.started, "Call init_broker() for proper initialization" # TODO message validation? try: message = json.dumps(message) if timeout: await wait_for( self.producer.send(topic=channel_name.lower(), value=message.encode(), **kwargs), timeout) else: await self.producer.send_and_wait(topic=channel_name.lower(), value=message.encode(), **kwargs) except Exception as e: raise SendMessageError(e) async def read(self, channel_name: List[str], timeout: int, **kwargs): raise NotImplementedError("Just a producer") async def close(self): await self.producer.stop()
class AIOKafkaRPCClient(object): log = logging.getLogger(__name__) def __init__(self, kafka_servers='localhost:9092', in_topic='aiokafkarpc_in', out_topic='aiokafkarpc_out', out_partitions=(0, ), translation_table=[], *, loop): self.call = CallObj(self._call_wrapper) self._topic_in = in_topic self._loop = loop self._waiters = {} self._out_topic = out_topic self._out_partitions = out_partitions default, ext_hook = get_msgpack_hooks(translation_table) self.__consumer = AIOKafkaConsumer( self._out_topic, loop=loop, bootstrap_servers=kafka_servers, group_id=None, key_deserializer=lambda x: x.decode("utf-8"), value_deserializer=lambda x: msgpack.unpackb( x, ext_hook=ext_hook, encoding="utf-8")) self.__producer = AIOKafkaProducer( bootstrap_servers=kafka_servers, loop=loop, key_serializer=lambda x: x.encode("utf-8"), value_serializer=lambda x: msgpack.packb(x, default=default)) @asyncio.coroutine def run(self): yield from self.__producer.start() yield from self.__consumer.start() # FIXME manual partition assignment does not work correctly in aiokafka # self.__consumer.assign( # [TopicPartition(self._out_topic, p) for p in self._out_partitions]) # # ensure that topic partitions exists for tp in self.__consumer.assignment(): yield from self.__consumer.position(tp) self._consume_task = self._loop.create_task(self.__consume_routine()) @asyncio.coroutine def close(self, timeout=10): yield from self.__producer.stop() if self._waiters: yield from asyncio.wait(self._waiters.values(), loop=self._loop, timeout=timeout) self._consume_task.cancel() try: yield from self._consume_task except asyncio.CancelledError: pass yield from self.__consumer.stop() for fut in self._waiters.values(): fut.set_exception(asyncio.TimeoutError()) def _call_wrapper(self, method): @asyncio.coroutine def rpc_call(*args, **kw_args): call_id = uuid.uuid4().hex ptid = random.choice(self._out_partitions) request = (method, args, kw_args, ptid) fut = asyncio.Future(loop=self._loop) fut.add_done_callback(lambda fut: self._waiters.pop(call_id)) self._waiters[call_id] = fut try: yield from self.__producer.send(self._topic_in, request, key=call_id) except Exception as err: self.log.error("send RPC request failed: %s", err) self._waiters[call_id].set_exception(err) return (yield from self._waiters[call_id]) return rpc_call @asyncio.coroutine def __consume_routine(self): while True: message = yield from self.__consumer.getone() call_id = message.key response = message.value self.call = CallObj(self._call_wrapper) fut = self._waiters.get(call_id) if fut is None: continue if "error" in response: self.log.debug(response.get("stacktrace")) fut.set_exception(RPCError(response["error"])) else: fut.set_result(response["result"])
class AIOKafkaRPC(object): log = logging.getLogger(__name__) def __init__(self, rpc_obj, kafka_servers='localhost:9092', in_topic='aiokafkarpc_in', out_topic='aiokafkarpc_out', translation_table=[], *, loop): self._tasks = {} self._loop = loop self._topic_out = out_topic self._rpc_obj = rpc_obj self._res_queue = asyncio.Queue(loop=loop) default, ext_hook = get_msgpack_hooks(translation_table) self.__consumer = AIOKafkaConsumer( in_topic, loop=loop, bootstrap_servers=kafka_servers, group_id=in_topic + '-group', key_deserializer=lambda x: x.decode("utf-8"), value_deserializer=lambda x: msgpack.unpackb( x, ext_hook=ext_hook, encoding="utf-8")) self.__producer = AIOKafkaProducer( bootstrap_servers=kafka_servers, loop=loop, key_serializer=lambda x: x.encode("utf-8"), value_serializer=lambda x: msgpack.packb(x, default=default)) @asyncio.coroutine def run(self): yield from self.__producer.start() yield from self.__consumer.start() self._consume_task = self._loop.create_task(self.__consume_routine()) self._produce_task = self._loop.create_task(self.__produce_routine()) @asyncio.coroutine def close(self, timeout=10): self._consume_task.cancel() try: yield from self._consume_task except asyncio.CancelledError: pass if self._tasks: yield from asyncio.wait(self._tasks.values(), loop=self._loop, timeout=timeout) self._res_queue.put_nowait((None, None, None)) yield from self._produce_task yield from self.__producer.stop() yield from self.__consumer.stop() def __send_result(self, call_id, ptid, result=None, err=None): if err is not None: out = io.StringIO() traceback.print_exc(file=out) ret = {"error": repr(err), "stacktrace": out.getvalue()} else: ret = {"result": result} self._res_queue.put_nowait((call_id, ptid, ret)) def __on_done_task(self, call_id, ptid, task): self._tasks.pop(call_id) try: result = task.result() self.__send_result(call_id, ptid, result) except Exception as err: self.__send_result(call_id, ptid, err=err) @asyncio.coroutine def __produce_routine(self): while True: call_id, ptid, res = yield from self._res_queue.get() if call_id is None: break try: yield from self.__producer.send(self._topic_out, res, key=call_id, partition=ptid) except KafkaError as err: self.log.error("send RPC response failed: %s", err) except Exception as err: self.__send_result(call_id, err=err) @asyncio.coroutine def __consume_routine(self): while True: message = yield from self.__consumer.getone() call_id = message.key method_name, args, kw_args, ptid = message.value try: method = getattr(self._rpc_obj, method_name) if not asyncio.iscoroutinefunction(method): raise RuntimeError( "'{}' should be a coroutine".format(method_name)) task = self._loop.create_task(method(*args, **kw_args)) task.add_done_callback( functools.partial(self.__on_done_task, call_id, ptid)) self._tasks[call_id] = task except Exception as err: self.__send_result(call_id, ptid, err=err)