Gracefully close Redpanda producer and consumer in example

Author: guilleiguaran

examples/redpanda_rspec.rb

@@ -12,6 +12,7 @@
 require "rdkafka"
 require "rspec"
 require "rspec/autorun"
+require "timeout"
 
 RSpec.configure do |config|
   config.add_setting :redpanda_container, default: nil
@@ -28,21 +29,44 @@
 
 RSpec.describe "Redpanda" do
   it "works" do
-    config = {
+    topic = "ruby-test-topic"
+    producer_config = {
+      "bootstrap.servers": RSpec.configuration.redpanda_container.connection_url,
+      "message.timeout.ms": 10_000
+    }
+    consumer_config = {
       "bootstrap.servers": RSpec.configuration.redpanda_container.connection_url,
       "group.id": "ruby-test",
       "auto.offset.reset": "earliest"
     }
 
-    producer = Rdkafka::Config.new(config).producer
-    producer.produce(payload: "test", topic: "ruby-test-topic").wait
+    producer = nil
+    consumer = nil
+
+    begin
+      producer = Rdkafka::Config.new(producer_config).producer
+      consumer = Rdkafka::Config.new(consumer_config).consumer
+      consumer.subscribe(topic)
+
+      producer.produce(payload: "test", topic: topic).wait
 
-    consumer = Rdkafka::Config.new(config).consumer
-    consumer.subscribe("ruby-test-topic")
+      message = nil
+      begin
+        Timeout.timeout(15) do
+          loop do
+            message = consumer.poll(1000)
+            break if message
+          end
+        end
+      rescue Timeout::Error
+        fail "Timed out waiting for message on #{topic}"
+      end
 
-    consumer.each do |message|
       expect(message.payload).to eq("test")
-      break
+    ensure
+      producer&.flush(5_000)
+      producer&.close
+      consumer&.close
     end
   end
 end