def test_message_queue_integration_enqueue_single_message(self):
"""Test that Queue enqueues single message properly"""
tq = Queue(self.config)
message = {"key": "value"}
tq.enqueue(message)
tq.engine.producer.close()
self.assertEqual(get_kafka_end_offset(tq, self.queue_name), 1)
def test_cli_integration_enqueues_raw(self):
"""Test that CLI enqueues raw messages from CLI options"""
tq = Queue(self.config)
self.assertEqual(get_kafka_end_offset(tq, self.queue_name), 0)
runner = CliRunner()
with runner.isolated_filesystem():
with open("raw_input.txt", "w") as f:
f.write(
'{"raw-key-1": "raw-value-1"}\n{"raw-key-2": "raw-value-2"}\n'
)
result = runner.invoke(
cli.main,
[
"--broker-url",
self.broker_url,
"--queue-name",
self.queue_name,
"enqueue-raw",
"--message-input",
"raw_input.txt",
],
obj={},
)
self.assertTrue(result.exit_code == 0)
self.assertEqual(get_kafka_end_offset(tq, self.queue_name), 2)
result = tq.dequeue()
self.assertEqual(result, b'{"raw-key-1": "raw-value-1"}')
result = tq.dequeue()
self.assertEqual(result, b'{"raw-key-2": "raw-value-2"}')
def setUp(self):
self.broker_url = "memory://"
self.queue_name = "test-message-queue"
self.config = deepcopy(Queue.get_default_config())
self.config["BROKER_URL"] = self.broker_url
self.config["QUEUE_NAME"] = self.queue_name
self.tq = Queue(self.config)
def test_message_queue_unit_enqueue(self, mock_connection):
"""Test that Queue mock enqueues message"""
message = {"key": "value"}
tq = Queue(self.config)
tq.enqueue(message)
tq.engine.queue.put.assert_called_once_with(message)
def test_message_queue_integration_connection(self):
"""Test that Queue connects properly"""
with closing(Queue(self.config)) as tq:
self.assertTrue(tq.engine.conn.connected)
self.assertFalse(tq.engine.conn.connected)
def enqueue(ctx: click.core.Context, message_entry: Tuple[Tuple[str, str]]):
"""Enqueue a message in the key/value format to the queue"""
# Convert input key/value to message
message: Dict[str, str] = {k: v for (k, v) in message_entry}
with closing(Queue(ctx.obj)) as queue:
queue.enqueue(message)
def setUp(self):
self.broker_url = CONFIG_TEST_IN_MEMORY_URL
self.queue_name = "test-message-queue"
self.config = deepcopy(Queue.get_default_config())
self.config["BROKER_URL"] = self.broker_url
self.config["QUEUE_NAME"] = self.queue_name
self.tq = Queue(self.config)
def test_cli_integration_enqueues_message_multiple_keys(self):
"""Test that CLI enqueues message from CLI options"""
tq = Queue(self.config)
self.assertEqual(get_kafka_end_offset(tq, self.queue_name), 0)
runner = CliRunner()
result = runner.invoke(
cli.main,
[
"--broker-url",
self.broker_url,
"--queue-name",
self.queue_name,
"enqueue",
"--message-entry",
"k1",
"v1",
"--message-entry",
"k2",
"v2",
"--message-entry",
"k3",
"v3",
],
obj={},
)
self.assertTrue(result.exit_code == 0)
self.assertEqual(get_kafka_end_offset(tq, self.queue_name), 1)
result = tq.dequeue()
self.assertEqual(result, b'{"k1": "v1", "k2": "v2", "k3": "v3"}')
def test_message_queue_unit_dequeue_empty(self, mock_producer,
mock_consumer):
"""Test that Queue mock dequeues empty message"""
tq = Queue(self.config)
tq.engine.consumer.__next__.side_effect = StopIteration()
self.assertEqual(tq.dequeue(), "{}")
def test_cli_integration_dequeues_batch(self):
"""Test that CLI dequeues batched message from CLI"""
tq = Queue(self.config)
self.assertEqual(get_kafka_end_offset(tq, self.queue_name), 0)
for i in range(3):
message = {f"cli-input-key-{i}": f"cli-input-value-{i}"}
tq.enqueue(message)
tq.engine.producer.close()
self.assertEqual(get_kafka_end_offset(tq, self.queue_name), 3)
runner = CliRunner()
result = runner.invoke(
cli.main,
[
"--broker-url",
self.broker_url,
"--queue-name",
self.queue_name,
"dequeue",
"--batch",
"3",
],
obj={},
)
self.assertTrue(result.exit_code == 0)
expected_messages = [
'{"cli-input-key-0": "cli-input-value-0"}\n',
'{"cli-input-key-1": "cli-input-value-1"}\n',
'{"cli-input-key-2": "cli-input-value-2"}\n',
]
self.assertEqual(result.output, "".join(expected_messages))
def test_message_queue_unit_enqueue(self, mock_producer, mock_consumer):
"""Test that Queue mock enqueues message"""
message = {"key": "value"}
tq = Queue(self.config)
tq.enqueue(message)
tq.engine.producer.send.assert_called_once_with(self.queue_name,
value=message)
def test_message_queue_unit_dequeue_raises_empty(self, mock_producer,
mock_consumer):
"""Test that Queue mock dequeues raises on empty message"""
self.config["RAISE_ON_EMPTY_DEQUEUE"] = True
tq = Queue(self.config)
tq.engine.consumer.__next__.side_effect = EmptyQueueException()
with self.assertRaises(EmptyQueueException):
tq.dequeue()
def test_message_queue_unit_dequeue(self, mock_producer, mock_consumer):
"""Test that Queue mock dequeues message"""
tq = Queue(self.config)
message = {"key": "value"}
msg = mock.Mock()
msg.value = message
tq.engine.consumer.__next__.return_value = msg
result = tq.dequeue()
tq.engine.consumer.__next__.assert_called_once()
self.assertDictEqual(result, message)
def test_message_queue_unit_close(self, mock_producer, mock_consumer):
"""Test that Queue connections are closed properly"""
with closing(Queue(self.config)) as tq:
tq.engine.consumer.commit.assert_not_called()
tq.engine.consumer.close.assert_not_called()
tq.engine.producer.close.assert_not_called()
tq.engine.consumer.commit.assert_called_once()
tq.engine.consumer.close.assert_called_once()
tq.engine.producer.close.assert_called_once()
def test_message_queue_unit_close(self, mock_connection):
"""Test that Queue connections are closed properly"""
mock_conn_obj = mock.Mock()
mock_connection.return_value = mock_conn_obj
with closing(Queue(self.config)) as tq:
tq.engine.queue.close.assert_not_called()
tq.engine.conn.release.assert_not_called()
tq.engine.queue.close.assert_called_once()
tq.engine.conn.release.assert_called_once()
def test_message_queue_integration_dequeue_single_message(self):
"""Test that Queue dequeues single message properly"""
tq = Queue(self.config)
self.assertEqual(get_kafka_end_offset(tq, self.queue_name), 0)
message = {"key-kafka-dequeue": "value"}
tq.enqueue(message)
tq.engine.producer.close()
self.assertEqual(get_kafka_end_offset(tq, self.queue_name), 1)
result = tq.dequeue()
tq.engine.consumer.commit()
self.assertEqual(result, b'{"key-kafka-dequeue": "value"}')
self.assertEqual(get_kafka_committed_offset(tq, self.queue_name), 1)
def test_message_queue_unit_dequeue(self, mock_connection):
"""Test that Queue mock dequeues message"""
msg = mock.Mock()
msg.body = b'{"key-kombu-dequeue": "value"}'
tq = Queue(self.config)
tq.engine.queue.get.return_value = msg
result = tq.dequeue()
tq.engine.queue.get.assert_called_once_with(
block=True,
timeout=self.config["DEQUEUE_TIMEOUT"],
)
self.assertEqual(result, b'{"key-kombu-dequeue": "value"}')
def test_message_queue_unit_init(self, mock_connection):
"""Test that Queue attributes are set properly"""
mock_conn_obj = mock.Mock()
mock_connection.return_value = mock_conn_obj
tq = Queue(self.config)
self.assertEqual(tq.engine.name, self.queue_name)
self.assertEqual(tq.engine.broker_url, self.broker_url)
mock_connection.assert_called_once_with(
self.broker_url, connect_timeout=self.config["CONNECTION_TIMEOUT"])
mock_conn_obj.SimpleQueue.assert_called_once_with(self.queue_name,
serializer="json")
def enqueue_raw(
ctx: click.core.Context,
message_input: TextIOWrapper,
raise_on_serialization_error: bool,
):
"""Enqueue raw json formatted messages to the queue"""
with closing(Queue(ctx.obj)) as queue:
for line in message_input:
try:
message = json.loads(line)
queue.enqueue(message)
except JSONDecodeError as exc:
if raise_on_serialization_error:
raise exc
click.echo(f"Cannot decode JSON: {repr(line)}")
click.echo(exc, err=True)
def test_message_queue_integration_dequeue_multiple_messages(self):
"""Test that Queue dequeues multiple message properly"""
tq = Queue(self.config)
self.assertEqual(get_kafka_end_offset(tq, self.queue_name), 0)
messages = [
{
"key1": "value1"
},
{
"key2": "value2"
},
{
"key3": "value3"
},
{
"key4": "value4"
},
]
for message in messages:
tq.enqueue(message)
tq.engine.producer.close()
self.assertEqual(get_kafka_end_offset(tq, self.queue_name), 4)
result = tq.dequeue()
tq.engine.consumer.commit()
self.assertEqual(result, b'{"key1": "value1"}')
self.assertEqual(get_kafka_committed_offset(tq, self.queue_name), 1)
result = tq.dequeue()
tq.engine.consumer.commit()
self.assertEqual(result, b'{"key2": "value2"}')
self.assertEqual(get_kafka_committed_offset(tq, self.queue_name), 2)
result = tq.dequeue()
tq.engine.consumer.commit()
self.assertEqual(result, b'{"key3": "value3"}')
self.assertEqual(get_kafka_committed_offset(tq, self.queue_name), 3)
result = tq.dequeue()
tq.engine.consumer.commit()
self.assertEqual(result, b'{"key4": "value4"}')
self.assertEqual(get_kafka_committed_offset(tq, self.queue_name), 4)
def test_message_queue_unit_init(self, mock_producer, mock_consumer):
"""Test that Queue attributes are set properly"""
tq = Queue(self.config)
self.assertEqual(tq.engine.topic, self.queue_name)
self.assertListEqual(tq.engine.servers,
["kafka_hostname1", "kafka_hostname2"])
mock_producer.assert_called_once_with(
bootstrap_servers=["kafka_hostname1", "kafka_hostname2"],
value_serializer=KafkaEngine.serializer,
)
mock_consumer.assert_called_once_with(
self.queue_name,
bootstrap_servers=["kafka_hostname1", "kafka_hostname2"],
auto_offset_reset=self.config["CONSUMER_AUTO_OFFSET_RESET"],
enable_auto_commit=self.config["CONSUMER_AUTOCOMMIT"],
group_id="poppy-test-message-queue",
consumer_timeout_ms=float("inf"),
)
def test_cli_integration_dequeues_empty(self):
"""Test that CLI dequeues empty message from CLI"""
tq = Queue(self.config)
self.assertEqual(get_kafka_end_offset(tq, self.queue_name), 0)
runner = CliRunner()
result = runner.invoke(
cli.main,
[
"--broker-url",
self.broker_url,
"--queue-name",
self.queue_name,
"dequeue",
"--blocking-dequeue-timeout",
"10",
],
obj={},
)
self.assertTrue(result.exit_code == 0)
def dequeue(
ctx: click.core.Context,
batch: int,
exit_on_empty: bool,
blocking_dequeue_timeout: int,
dequeue_raise_on_empty: bool,
consumer_group_id: str,
consumer_autocommit: bool,
consumer_auto_offset_reset: str,
):
"""Dequeue message from the queue"""
ctx.obj["DEQUEUE_TIMEOUT"] = blocking_dequeue_timeout
ctx.obj["RAISE_ON_EMPTY_DEQUEUE"] = dequeue_raise_on_empty
ctx.obj["CONSUMER_AUTOCOMMIT"] = consumer_autocommit
ctx.obj["CONSUMER_AUTO_OFFSET_RESET"] = consumer_auto_offset_reset
ctx.obj["DEQUEUE_EXIT_ON_EMPTY"] = exit_on_empty
if consumer_group_id:
ctx.obj["CONSUMER_GROUP_ID"] = consumer_group_id
if exit_on_empty and not blocking_dequeue_timeout:
msg = "--exit-on-empty needs to be combined with --blocking-dequeue-timeout argument"
raise click.exceptions.ClickException(msg)
if exit_on_empty and not dequeue_raise_on_empty:
msg = "--exit-on-empty needs to be combined with --dequeue-raise-on-empty=True"
raise click.exceptions.ClickException(msg)
with closing(Queue(ctx.obj)) as queue:
for _ in range(batch):
try:
message = queue.dequeue()
click.echo(message)
except EmptyQueueException as exc:
if exit_on_empty:
msg = "Queue is empty"
click.echo(msg, err=True)
sys.exit(100)
raise exc
def test_cli_integration_dequeues_message(self):
"""Test that CLI dequeues message from CLI"""
tq = Queue(self.config)
message = {"cli-input-key": "cli-input-value"}
self.assertEqual(get_kafka_end_offset(tq, self.queue_name), 0)
tq.enqueue(message)
tq.engine.producer.close()
self.assertEqual(get_kafka_end_offset(tq, self.queue_name), 1)
runner = CliRunner()
result = runner.invoke(
cli.main,
[
"--broker-url",
self.broker_url,
"--queue-name",
self.queue_name,
"dequeue",
],
obj={},
)
self.assertTrue(result.exit_code == 0)
self.assertEqual({"cli-input-key": "cli-input-value"}, message)
def test_message_queue_integration_enqueue_multiple_messages(self):
"""Test that Queue enqueues multiple messages properly"""
tq = Queue(self.config)
self.assertEqual(get_kafka_end_offset(tq, self.queue_name), 0)
messages = [
{
"key1": "value1"
},
{
"key2": "value2"
},
{
"key3": "value3"
},
{
"key4": "value4"
},
]
for message in messages:
tq.enqueue(message)
tq.engine.producer.close()
self.assertEqual(get_kafka_end_offset(tq, self.queue_name), 4)