from tests.pykafka import test_simpleconsumer
from pykafka.rdkafka import RdKafkaSimpleConsumer


class TestRdKafkaSimpleConsumer(test_simpleconsumer.TestSimpleConsumer):
    def _get_simple_consumer(self, **kwargs):
        # This enables automatic reuse of all tests from test_simpleconsumer
        topic = self.client.topics[self.topic_name]
        return RdKafkaSimpleConsumer(
            topic=topic,
            cluster=topic._cluster,
            **kwargs)

    def test_offset_commit_agrees(self):
        """Check rdkafka-obtained offsets arrive correctly

        In RdKafkaSimpleConsumer.consume we bypass most of the internals of
        simpleconsumer.OwnedPartition, but then expect it to still commit
        offsets for us correctly.  This warrants very explicit testing.
        """
        consumer = self._get_simple_consumer(
            consumer_group='test_offset_commit_agrees')
        latest_offs = _latest_partition_offsets_by_reading(consumer, 100)

        consumer.commit_offsets()
        res = consumer.fetch_offsets()
        consumer.stop()

        # Discard any empty partitions (which thus won't be in latest_offs):
        retrieved_offs = {r[0]: r[1].offset for r in res if r[1].offset != -1}
        self.assertEquals(retrieved_offs, latest_offs)

    def test_offset_resume_agrees(self):
        """Check the rdkafka consumer returns messages at specified offset

        Make sure reads from the underlying rdkafka consumer really do start
        at the offsets dictated by SimpleConsumer
        """
        consumer = self._get_simple_consumer(
            consumer_group='test_offset_resume_agrees')
        latest_offs = _latest_partition_offsets_by_reading(consumer, 100)
        consumer.commit_offsets()
        consumer.stop()

        consumer = self._get_simple_consumer(
            consumer_group='test_offset_resume_agrees')
        while latest_offs: # check each partition, then tick it off:
            msg = consumer.consume()
            expected_offset = latest_offs[msg.partition_id] + 1
            self.assertEquals(msg.offset, expected_offset)
            del latest_offs[msg.partition_id]
        consumer.stop()


def _latest_partition_offsets_by_reading(consumer, n_reads):
    """Obtain message offsets from consumer, return grouped by partition"""
    latest_offs = {}
    for _ in xrange(n_reads):
        msg = consumer.consume()
        latest_offs[msg.partition_id] = msg.offset
    return latest_offs

@@ -972,17 +972,16 @@
 r, 100)%0A
-%0A
         
@@ -1006,24 +1006,25 @@
 fsets()%0A
+%0A
         
 res = co
@@ -1019,219 +1019,201 @@
     
-res = consumer.fetch_offsets()%0A        consumer.stop()%0A%0A        # Discard any empty partitions (which thus won't be in latest_offs):%0A        retrieved_offs = %7Br%5B0%5D: r%5B1%5D.offset for r in res if r%5B1%5D.offset != -1%7D
+# We can only compare partitions we've consumed from, so filter:%0A        retrieved_offs = %7Br%5B0%5D: r%5B1%5D.offset for r in consumer.fetch_offsets()%0A                          if r%5B0%5D in latest_offs%7D%0A
 %0A   
@@ -1263,16 +1263,40 @@
 st_offs)
+%0A        consumer.stop()
 %0A%0A    de
@@ -1980,32 +1980,173 @@
 sumer.consume()%0A
+            if msg.partition_id not in latest_offs:%0A                # ie we didn't get to this partition previously%0A                continue%0A
             expe

fix offset tests