Handle assignment/revocation in consumer mock

mfontanini · mfontanini · commit 798254d2c58f · 2017-10-21T17:17:34.000-07:00
diff --git a/mocking/include/cppkafka/mocking/consumer_mock.h b/mocking/include/cppkafka/mocking/consumer_mock.h
@@ -4,6 +4,8 @@
 #include <vector>
 #include <string>
 #include <cstdint>
+#include <map>
+#include <tuple>
 #include <unordered_set>
 #include <cppkafka/mocking/handle_mock.h>
 #include <cppkafka/mocking/configuration_mock.h>
@@ -20,16 +22,31 @@ class ConsumerMock : public HandleMock {
     ~ConsumerMock();
 
     void subscribe(const std::vector<std::string>& topics);
+    void set_opaque(void* opaque);
 private:
     static uint64_t make_consumer_id();
 
+    struct TopicPartitionInfo {
+        int64_t offset;
+        bool paused;
+    };
+
+    using TopicPartitionId = std::tuple<std::string, int>;
+
+    static TopicPartitionId make_id(const TopicPartitionMock& topic_partition);
     void on_assignment(std::vector<TopicPartitionMock>& topic_partitions);
     void on_revocation(const std::vector<TopicPartitionMock>& topic_partitions);
     void on_message(uint64_t offset);
+    template <typename List>
+    void handle_rebalance(rd_kafka_resp_err_t type, List& topic_partitions);
+    void handle_assign(const TopicPartitionMock& topic_partition);
+    void handle_unassign(const TopicPartitionMock& topic_partition);
 
     ConfigurationMock config_;
     const std::string group_id_;
     std::unordered_set<std::string> subscribed_topics_;
+    std::map<TopicPartitionId, TopicPartitionInfo> assigned_partitions_;
+    void* opaque_;
     uint64_t consumer_id_;
 };
 
diff --git a/mocking/include/cppkafka/mocking/topic_partition_mock.h b/mocking/include/cppkafka/mocking/topic_partition_mock.h
@@ -3,6 +3,8 @@
 
 #include <string>
 #include <cstdint>
+#include <memory>
+#include <vector>
 #include <librdkafka/rdkafka.h>
 
 namespace cppkafka {
@@ -18,11 +20,18 @@ class TopicPartitionMock {
 
     void set_offset(int64_t offset);
 private:
-    std::string topic_;
-    int partition_;
+    const std::string topic_;
+    const int partition_;
     int64_t offset_;
 };
 
+using TopicPartitionMockListPtr = std::unique_ptr<rd_kafka_topic_partition_list_t,
+                                                 decltype(&rd_kafka_topic_partition_list_destroy)>;
+TopicPartitionMockListPtr
+to_rdkafka_handle(const std::vector<TopicPartitionMock>& topic_partitions);
+std::vector<TopicPartitionMock>
+from_rdkafka_handle(const rd_kafka_topic_partition_list_t& topic_partitions);
+
 } // mocking
 } // cppkafka
 
diff --git a/mocking/src/api.cpp b/mocking/src/api.cpp
@@ -201,6 +201,7 @@ rd_kafka_topic_partition_list_add(rd_kafka_topic_partition_list_t* toppar_list,
     output->topic = new char[length + 1];
     copy(topic, topic + length, output->topic);
     output->partition = partition;
+    output->offset = RD_KAFKA_OFFSET_INVALID;
     return output;
 }
 
diff --git a/mocking/src/consumer_mock.cpp b/mocking/src/consumer_mock.cpp
@@ -9,6 +9,7 @@ using std::string;
 using std::move;
 using std::bind;
 using std::runtime_error;
+using std::make_tuple;
 
 namespace cppkafka {
 namespace mocking {
@@ -57,16 +58,53 @@ void ConsumerMock::subscribe(const vector<string>& topics) {
     }
 }
 
-void ConsumerMock::on_assignment(vector<TopicPartitionMock>& topic_partitions) {
+void ConsumerMock::set_opaque(void* opaque) {
+    opaque_ = opaque;
+}
 
+ConsumerMock::TopicPartitionId ConsumerMock::make_id(const TopicPartitionMock& topic_partition) {
+    return make_tuple(topic_partition.get_topic(), topic_partition.get_partition());
 }
 
-void ConsumerMock::on_revocation(const vector<TopicPartitionMock>& topic_partitions) {
+void ConsumerMock::on_assignment(vector<TopicPartitionMock>& topic_partitions) {
+    handle_rebalance(RD_KAFKA_RESP_ERR__ASSIGN_PARTITIONS, topic_partitions);
+    for (const TopicPartitionMock& topic_partition : topic_partitions) {
+        handle_assign(topic_partition);
+    }
+}
 
+void ConsumerMock::on_revocation(const vector<TopicPartitionMock>& topic_partitions) {
+    handle_rebalance(RD_KAFKA_RESP_ERR__REVOKE_PARTITIONS, topic_partitions);
+    for (const TopicPartitionMock& topic_partition : topic_partitions) {
+        handle_unassign(topic_partition);
+    }
 }
 
 void ConsumerMock::on_message(uint64_t offset) {
+    
+}
+
+template <typename List>
+void ConsumerMock::handle_rebalance(rd_kafka_resp_err_t type, List& topic_partitions) {
+    auto rebalance_callback = config_.get_rebalance_callback();
+    if (rebalance_callback) {
+        auto handle = to_rdkafka_handle(topic_partitions);
+        rebalance_callback(nullptr, type, handle.get(), opaque_);
+    }
+}
+
+void ConsumerMock::handle_assign(const TopicPartitionMock& topic_partition) {
+    const auto id = make_id(topic_partition);
+    if (assigned_partitions_.count(id)) {
+        return;
+    }
+    assigned_partitions_[id] = {
+        topic_partition.get_offset()
+    };
+}
 
+void ConsumerMock::handle_unassign(const TopicPartitionMock& topic_partition) {
+    assigned_partitions_.erase(make_id(topic_partition));
 }
 
 } // mocking
diff --git a/mocking/src/topic_partition_mock.cpp b/mocking/src/topic_partition_mock.cpp
@@ -1,6 +1,10 @@
+#include <tuple>
 #include <cppkafka/mocking/topic_partition_mock.h>
 
 using std::string;
+using std::vector;
+using std::tie;
+using std::unique_ptr;
 
 namespace cppkafka {
 namespace mocking {
@@ -26,5 +30,29 @@ void TopicPartitionMock::set_offset(int64_t offset) {
     offset_ = offset;
 }
 
+TopicPartitionMockListPtr to_rdkafka_handle(const vector<TopicPartitionMock>& topic_partitions){
+    const size_t count = topic_partitions.size();
+    TopicPartitionMockListPtr output{rd_kafka_topic_partition_list_new(count),
+                                     &rd_kafka_topic_partition_list_destroy};
+    for (const TopicPartitionMock& topic_partition : topic_partitions) {
+        auto* ptr = rd_kafka_topic_partition_list_add(output.get(),
+                                                      topic_partition.get_topic().data(),
+                                                      topic_partition.get_partition());
+        ptr->offset = topic_partition.get_offset();
+    }
+    return output;
+}
+
+vector<TopicPartitionMock>
+from_rdkafka_handle(const rd_kafka_topic_partition_list_t& topic_partitions){
+    vector<TopicPartitionMock> output;
+    for (int i = 0; i < topic_partitions.cnt; ++i) {
+        const auto& topic_partition = topic_partitions.elems[i];
+        output.emplace_back(topic_partition.topic, topic_partition.partition,
+                            topic_partition.offset);
+    }
+    return output;
+}
+
 } // mocking
 } // cppkafka

Original file line number	Diff line number	Diff line change
`@@ -201,6 +201,7 @@ rd_kafka_topic_partition_list_add(rd_kafka_topic_partition_list_t* toppar_list,`
`201`	`201`	`output->topic = new char[length + 1];`
`202`	`202`	`copy(topic, topic + length, output->topic);`
`203`	`203`	`output->partition = partition;`
	`204`	`+ output->offset = RD_KAFKA_OFFSET_INVALID;`
`204`	`205`	`return output;`
`205`	`206`	`}`
`206`	`207`