finish tests

Author: rerpha

tests/test_consume.py

@@ -82,3 +82,15 @@ def test_consume_but_exception_thrown_consumer_is_closed():
         c.return_value.consume.side_effect = Exception
         consume("somebroker", "sometopic", num_messages=1)
         c.return_value.close.assert_called_once()
+
+@patch("saluki.consume.Consumer")
+def test_consume_with_timestamp(mock_consumer):
+    expected_topic = "sometopic"
+    partition = 0
+    timestamp = 1234
+    offset = 2345
+
+    mock_consumer.offsets_for_times.return_value = [TopicPartition(expected_topic, partition, offset)]
+    consume("somebroker", topic=expected_topic, timestamp=timestamp, partition=partition)
+
+    mock_consumer.return_value.assign.assert_called_with([TopicPartition(expected_topic, partition, offset)])

tests/test_play.py

@@ -37,10 +37,10 @@ def test_play_with_offsets():
         consumer_obj.consume.assert_called_with(2)  # stop - start + 1
 
         p_obj = p()
-        call_1 = p_obj.produce.call_args_list[0]
-        assert call_1.args == (dest_topic, message_1_val, message_1_key)
-        call_2 = p_obj.produce.call_args_list[1]
-        assert call_2.args == (dest_topic, message_2_val, message_2_key)
+        produce_batch_call = p_obj.produce_batch.call_args.args
+        assert dest_topic == produce_batch_call[0]
+        assert {'key': message_1_key, 'value': message_1_val} in produce_batch_call[1]
+        assert {'key': message_2_key, 'value': message_2_val} in produce_batch_call[1]
 
 
 def test_play_with_timestamps():
@@ -64,9 +64,9 @@ def test_play_with_timestamps():
 
     with patch("saluki.play.Consumer") as c, patch("saluki.play.Producer") as p:
         consumer_obj = c()
-        consumer_obj.offsets_for_times.return_value = [
-            TopicPartition(src_topic, partition=0, offset=2),
-            TopicPartition(src_topic, partition=0, offset=3),
+        consumer_obj.offsets_for_times.side_effect = [
+            [TopicPartition(src_topic, partition=0, offset=2)],
+            [TopicPartition(src_topic, partition=0, offset=3)]
         ]
         consumer_obj.consume.return_value = [message_1, message_2]
 
@@ -78,10 +78,10 @@ def test_play_with_timestamps():
         consumer_obj.consume.assert_called_with(2)  # stop - start + 1
 
         p_obj = p()
-        call_1 = p_obj.produce.call_args_list[0]
-        assert call_1.args == (dest_topic, message_1_val, message_1_key)
-        call_2 = p_obj.produce.call_args_list[1]
-        assert call_2.args == (dest_topic, message_2_val, message_2_key)
+        produce_batch_call = p_obj.produce_batch.call_args.args
+        assert dest_topic == produce_batch_call[0]
+        assert {'key': message_1_key, 'value': message_1_val} in produce_batch_call[1]
+        assert {'key': message_2_key, 'value': message_2_val} in produce_batch_call[1]
 
 
 def test_play_with_exception_when_consuming_consumer_still_closed():

tests/test_utils.py

@@ -1,3 +1,4 @@
+from argparse import ArgumentTypeError
 from unittest.mock import Mock, patch
 
 import pytest
@@ -13,7 +14,7 @@
     _parse_timestamp,
     _try_to_deserialise_message,
     deserialise_and_print_messages,
-    parse_kafka_uri,
+    parse_kafka_uri, dateutil_parsable_or_unix_timestamp,
 )
 
 
@@ -174,3 +175,27 @@ def test_uri_with_no_topic():
     test_broker = "some_broker"
     with pytest.raises(RuntimeError):
         parse_kafka_uri(test_broker)
+
+@pytest.mark.parametrize("timestamp",
+["2025-11-19T15:27:11",
+ "2025-11-19T15:27:11Z",
+"2025-11-19T15:27:11+00:00"
+ ]
+
+)
+def test_parses_datetime_properly_with_string(timestamp):
+    assert dateutil_parsable_or_unix_timestamp(timestamp) == 1763566031000
+
+@pytest.mark.parametrize("timestamp",
+["1763566031000",
+ "1763566031",
+ "1763566031000000",
+ ]
+)
+def test_parses_datetime_properly_and_leaves_unix_timestamp_alone(timestamp):
+    assert dateutil_parsable_or_unix_timestamp(timestamp) == int(timestamp)
+
+
+def test_invalid_timestamp_raises():
+    with pytest.raises(ArgumentTypeError):
+        dateutil_parsable_or_unix_timestamp("invalid")