Fix timestamp not passed to RecordMetadata (#1273)

* Fix timestamp not being passed to RecordMetadata properly

* Add more tests for LegacyBatch

* Fix producer test for recordmetadata

2 files changed, 139 insertions, 8 deletions

diff --git a/test/record/test_legacy_records.py b/test/record/test_legacy_records.py
index 2d76695..ffe8a35 100644
--- a/test/record/test_legacy_records.py
+++ b/test/record/test_legacy_records.py
@@ -1,8 +1,8 @@
+from __future__ import unicode_literals
 import pytest
 from kafka.record.legacy_records import (
     LegacyRecordBatch, LegacyRecordBatchBuilder
 )
-from kafka.protocol.message import Message
 
 
 @pytest.mark.parametrize("magic", [0, 1])
@@ -27,9 +27,9 @@ def test_read_write_serde_v0_v1_no_compression(magic):
 
 
 @pytest.mark.parametrize("compression_type", [
-    Message.CODEC_GZIP,
-    Message.CODEC_SNAPPY,
-    Message.CODEC_LZ4
+    LegacyRecordBatch.CODEC_GZIP,
+    LegacyRecordBatch.CODEC_SNAPPY,
+    LegacyRecordBatch.CODEC_LZ4
 ])
 @pytest.mark.parametrize("magic", [0, 1])
 def test_read_write_serde_v0_v1_with_compression(compression_type, magic):
@@ -43,14 +43,14 @@ def test_read_write_serde_v0_v1_with_compression(compression_type, magic):
     batch = LegacyRecordBatch(bytes(buffer), magic)
     msgs = list(batch)
 
-    expected_checksum = (-2095076219 if magic else 278251978) & 0xffffffff
     for offset, msg in enumerate(msgs):
         assert msg.offset == offset
         assert msg.timestamp == (9999999 if magic else None)
         assert msg.timestamp_type == (0 if magic else None)
         assert msg.key == b"test"
         assert msg.value == b"Super"
-        assert msg.checksum == expected_checksum
+        assert msg.checksum == (-2095076219 if magic else 278251978) & \
+            0xffffffff
 
 
 @pytest.mark.parametrize("magic", [0, 1])
@@ -83,3 +83,84 @@ def test_estimate_size_in_bytes_bigger_than_batch(magic):
     buf = builder.build()
     assert len(buf) <= estimate_size, \
         "Estimate should always be upper bound"
+
+
+@pytest.mark.parametrize("magic", [0, 1])
+def test_legacy_batch_builder_validates_arguments(magic):
+    builder = LegacyRecordBatchBuilder(
+        magic=magic, compression_type=0, batch_size=1024 * 1024)
+
+    # Key should not be str
+    with pytest.raises(TypeError):
+        builder.append(
+            0, timestamp=9999999, key="some string", value=None)
+
+    # Value should not be str
+    with pytest.raises(TypeError):
+        builder.append(
+            0, timestamp=9999999, key=None, value="some string")
+
+    # Timestamp should be of proper type
+    if magic != 0:
+        with pytest.raises(TypeError):
+            builder.append(
+                0, timestamp="1243812793", key=None, value=b"some string")
+
+    # Offset of invalid type
+    with pytest.raises(TypeError):
+        builder.append(
+            "0", timestamp=9999999, key=None, value=b"some string")
+
+    # Ok to pass value as None
+    builder.append(
+        0, timestamp=9999999, key=b"123", value=None)
+
+    # Timestamp can be None
+    builder.append(
+        1, timestamp=None, key=None, value=b"some string")
+
+    # Ok to pass offsets in not incremental order. This should not happen thou
+    builder.append(
+        5, timestamp=9999999, key=b"123", value=None)
+
+    # in case error handling code fails to fix inner buffer in builder
+    assert len(builder.build()) == 119 if magic else 95
+
+
+@pytest.mark.parametrize("magic", [0, 1])
+def test_legacy_correct_metadata_response(magic):
+    builder = LegacyRecordBatchBuilder(
+        magic=magic, compression_type=0, batch_size=1024 * 1024)
+    meta = builder.append(
+        0, timestamp=9999999, key=b"test", value=b"Super")
+
+    assert meta.offset == 0
+    assert meta.timestamp == (9999999 if magic else -1)
+    assert meta.crc == (-2095076219 if magic else 278251978) & 0xffffffff
+    assert repr(meta) == (
+        "LegacyRecordMetadata(offset=0, crc={}, size={}, "
+        "timestamp={})".format(meta.crc, meta.size, meta.timestamp)
+    )
+
+
+@pytest.mark.parametrize("magic", [0, 1])
+def test_legacy_batch_size_limit(magic):
+    # First message can be added even if it's too big
+    builder = LegacyRecordBatchBuilder(
+        magic=magic, compression_type=0, batch_size=1024)
+    meta = builder.append(0, timestamp=None, key=None, value=b"M" * 2000)
+    assert meta.size > 0
+    assert meta.crc is not None
+    assert meta.offset == 0
+    assert meta.timestamp is not None
+    assert len(builder.build()) > 2000
+
+    builder = LegacyRecordBatchBuilder(
+        magic=magic, compression_type=0, batch_size=1024)
+    meta = builder.append(0, timestamp=None, key=None, value=b"M" * 700)
+    assert meta is not None
+    meta = builder.append(1, timestamp=None, key=None, value=b"M" * 700)
+    assert meta is None
+    meta = builder.append(2, timestamp=None, key=None, value=b"M" * 700)
+    assert meta is None
+    assert len(builder.build()) < 1000
diff --git a/test/test_producer.py b/test/test_producer.py
index 1f6608a..41bd52e 100644
--- a/test/test_producer.py
+++ b/test/test_producer.py
@@ -1,11 +1,11 @@
 import gc
 import platform
-import sys
+import time
 import threading
 
 import pytest
 
-from kafka import KafkaConsumer, KafkaProducer
+from kafka import KafkaConsumer, KafkaProducer, TopicPartition
 from kafka.producer.buffer import SimpleBufferPool
 from test.conftest import version
 from test.testutil import random_string
@@ -78,3 +78,53 @@ def test_kafka_producer_gc_cleanup():
     del(producer)
     gc.collect()
     assert threading.active_count() == threads
+
+
+@pytest.mark.skipif(not version(), reason="No KAFKA_VERSION set")
+@pytest.mark.parametrize("compression", [None, 'gzip', 'snappy', 'lz4'])
+def test_kafka_producer_proper_record_metadata(kafka_broker, compression):
+    connect_str = ':'.join([kafka_broker.host, str(kafka_broker.port)])
+    producer = KafkaProducer(bootstrap_servers=connect_str,
+                             retries=5,
+                             max_block_ms=10000,
+                             compression_type=compression)
+    if producer.config['api_version'] >= (0, 10):
+        magic = 1
+    else:
+        magic = 0
+
+    topic = random_string(5)
+    future = producer.send(
+        topic,
+        value=b"Simple value", key=b"Simple key", timestamp_ms=9999999,
+        partition=0)
+    record = future.get(timeout=5)
+    assert record is not None
+    assert record.topic == topic
+    assert record.partition == 0
+    assert record.topic_partition == TopicPartition(topic, 0)
+    assert record.offset == 0
+    if magic >= 1:
+        assert record.timestamp == 9999999
+    else:
+        assert record.timestamp == -1  # NO_TIMESTAMP
+
+    if magic == 1:
+        assert record.checksum == 1370034956
+    else:
+        assert record.checksum == 3296137851
+
+    assert record.serialized_key_size == 10
+    assert record.serialized_value_size == 12
+
+    # generated timestamp case is skipped for broker 0.9 and below
+    if magic == 0:
+        return
+
+    send_time = time.time() * 1000
+    future = producer.send(
+        topic,
+        value=b"Simple value", key=b"Simple key", timestamp_ms=None,
+        partition=0)
+    record = future.get(timeout=5)
+    assert abs(record.timestamp - send_time) <= 1000  # Allow 1s deviation