Merge pull request #529 from zendesk/dasch/improve-consumer-specs

Improve the consumer specs

Author: dasch

spec/functional/consumer_group_spec.rb

@@ -1,102 +1,77 @@
 describe "Consumer API", functional: true do
-  let(:num_partitions) { 15 }
-  let!(:topic) { create_random_topic(num_partitions: 3) }
   let(:offset_retention_time) { 30 }
 
   example "consuming messages from the beginning of a topic" do
+    topic = create_random_topic(num_partitions: 1)
     messages = (1..1000).to_a
 
-    Thread.new do
+    begin
       kafka = Kafka.new(seed_brokers: kafka_brokers, client_id: "test")
       producer = kafka.producer
 
       messages.each do |i|
-        producer.produce(i.to_s, topic: topic, partition_key: i.to_s)
-
-        if i % 100 == 0
-          producer.deliver_messages
-          sleep 1
-        end
-      end
-
-      (0...num_partitions).each do |i|
-        # Send a tombstone to each partition.
-        producer.produce(nil, topic: topic, partition: i)
+        producer.produce(i.to_s, topic: topic, partition: 0)
       end
 
       producer.deliver_messages
     end
 
     group_id = "test#{rand(1000)}"
 
-    threads = 2.times.map do |thread_id|
-      t = Thread.new do
-        received_messages = []
+    mutex = Mutex.new
+    received_messages = []
 
-        kafka = Kafka.new(seed_brokers: kafka_brokers, client_id: "test", logger: logger)
-        consumer = kafka.consumer(group_id: group_id, offset_retention_time: offset_retention_time)
-        consumer.subscribe(topic)
+    consumers = 2.times.map do
+      kafka = Kafka.new(seed_brokers: kafka_brokers, client_id: "test", logger: logger)
+      consumer = kafka.consumer(group_id: group_id, offset_retention_time: offset_retention_time)
+      consumer.subscribe(topic)
+      consumer
+    end
 
+    threads = consumers.map do |consumer|
+      t = Thread.new do
         consumer.each_message do |message|
-          if message.value.nil?
-            consumer.stop
-          else
-            received_messages << Integer(message.value)
+          mutex.synchronize do
+            received_messages << message
+
+            if received_messages.count == messages.count
+              consumers.each(&:stop)
+            end
           end
         end
-
-        received_messages
       end
 
       t.abort_on_exception = true
 
       t
     end
 
-    received_messages = threads.map(&:value).flatten
+    threads.each(&:join)
 
-    expect(received_messages.sort).to match_array messages
+    expect(received_messages.map(&:value).map(&:to_i)).to match_array messages
   end
 
-  example "consuming messages from the end of a topic" do
-    sent_messages = 1_000
-
-    num_partitions = 1
+  example "consuming messages from a topic that's being written to" do
+    num_partitions = 3
     topic = create_random_topic(num_partitions: num_partitions)
-    group_id = "test#{rand(1000)}"
+    messages = (1..100).to_a
 
-    consumer_thread = Thread.new do
-      received_messages = 0
-
-      kafka = Kafka.new(seed_brokers: kafka_brokers, client_id: "test", logger: logger)
-      consumer = kafka.consumer(group_id: group_id, offset_retention_time: offset_retention_time)
-      consumer.subscribe(topic, start_from_beginning: false)
-
-      consumer.each_message do |message|
-        if message.value.nil?
-          consumer.stop
-        else
-          received_messages += 1
-        end
-      end
-
-      received_messages
-    end
-
-    consumer_thread.abort_on_exception = true
-
-    sleep 1
+    mutex = Mutex.new
+    var = ConditionVariable.new
 
     Thread.new do
       kafka = Kafka.new(seed_brokers: kafka_brokers, client_id: "test")
       producer = kafka.producer
 
-      1.upto(sent_messages) do |i|
-        producer.produce("hello", topic: topic, partition_key: i.to_s)
+      messages.each do |i|
+        producer.produce(i.to_s, topic: topic, partition: i % 3)
 
         if i % 100 == 0
           producer.deliver_messages
-          sleep 1
+
+          mutex.synchronize do
+            var.wait(mutex)
+          end
         end
       end
 
@@ -108,9 +83,35 @@
       producer.deliver_messages
     end
 
-    received_messages = consumer_thread.value
+    group_id = "test#{rand(1000)}"
+    received_messages = []
+
+    threads = 2.times.map do |thread_id|
+      t = Thread.new do
+        kafka = Kafka.new(seed_brokers: kafka_brokers, client_id: "test", logger: logger)
+        consumer = kafka.consumer(group_id: group_id, offset_retention_time: offset_retention_time)
+        consumer.subscribe(topic)
+
+        consumer.each_message do |message|
+          if message.value.nil?
+            consumer.stop
+          else
+            mutex.synchronize do
+              received_messages << Integer(message.value)
+              var.signal
+            end
+          end
+        end
+      end
+
+      t.abort_on_exception = true
+
+      t
+    end
+
+    threads.each(&:join)
 
-    expect(received_messages).to eq sent_messages
+    expect(received_messages).to match_array messages
   end
 
   example "stopping and restarting a consumer group" do