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)
