fix: test_produce_consume_base_assign

Author: j7nw4r

tests/admin.rs

@@ -4,13 +4,16 @@ use crate::utils::logging::init_test_logger;
 use crate::utils::rand::{rand_test_group, rand_test_topic};
 use crate::utils::{get_broker_version, KafkaVersion};
 use backon::{BlockingRetryable, ExponentialBuilder};
-use rdkafka::admin::{AdminClient, AdminOptions, AlterConfig, ConfigEntry, ConfigSource, GroupResult, NewPartitions, NewTopic, OwnedResourceSpecifier, ResourceSpecifier, TopicReplication};
+use rdkafka::admin::{
+    AdminClient, AdminOptions, AlterConfig, ConfigEntry, ConfigSource, GroupResult, NewPartitions,
+    NewTopic, OwnedResourceSpecifier, ResourceSpecifier, TopicReplication,
+};
+use rdkafka::client::DefaultClientContext;
 use rdkafka::error::KafkaError;
 use rdkafka::producer::{FutureRecord, Producer};
 use rdkafka::{ClientConfig, Offset, TopicPartitionList};
 use rdkafka_sys::RDKafkaErrorCode;
 use std::time::Duration;
-use rdkafka::client::DefaultClientContext;
 
 #[path = "utils/mod.rs"]
 mod utils;
@@ -666,4 +669,4 @@ async fn test_event_errors() {
         res,
         Err(KafkaError::AdminOp(RDKafkaErrorCode::OperationTimedOut))
     );
-}
+}

tests/test_high_consumers.rs renamed to tests/stream_consumers.rs

@@ -15,10 +15,14 @@ use rdkafka::error::KafkaError;
 use rdkafka::topic_partition_list::{Offset, TopicPartitionList};
 use rdkafka::util::current_time_millis;
 use rdkafka::{Message, Timestamp};
+use rdkafka::admin::AdminOptions;
 use rdkafka_sys::types::RDKafkaConfRes;
 
+use crate::utils::containers::KafkaContext;
+use crate::utils::logging::init_test_logger;
 use crate::utils::rand::*;
 use crate::utils::*;
+use crate::utils::admin::new_topic_vec;
 
 mod utils;
 
@@ -71,27 +75,52 @@ async fn test_invalid_max_poll_interval() {
 // All produced messages should be consumed.
 #[tokio::test(flavor = "multi_thread")]
 async fn test_produce_consume_base() {
-    let _r = env_logger::try_init();
+    init_test_logger();
+
+    // Get Kafka container context.
+    let kafka_context = KafkaContext::shared()
+        .await
+        .expect("could not create kafka context");
 
+    let producer = producer::future_producer::create_producer(&kafka_context.bootstrap_servers)
+        .await
+        .expect("Could not create Future producer");
+
+    let num_of_messages_to_send = 100usize;
     let start_time = current_time_millis();
     let topic_name = rand_test_topic("test_produce_consume_base");
-    let message_map = populate_topic(&topic_name, 100, &value_fn, &key_fn, None, None).await;
-    let consumer = create_stream_consumer(&rand_test_group(), None);
-    consumer.subscribe(&[topic_name.as_str()]).unwrap();
+    let message_map = topics::populate_topic_using_future_producer(
+        &producer,
+        &topic_name,
+        num_of_messages_to_send,
+        None,
+    )
+    .await
+    .expect("Could not populate topic using Future producer");
+    // let message_map = populate_topic(&topic_name, 100, &value_fn, &key_fn, None, None).await;
+    let consumer = utils::consumer::stream_consumer::create_stream_consumer(
+        &kafka_context.bootstrap_servers,
+        Some(&rand_test_group()),
+    )
+    .await
+    .expect("could not create stream consumer");
+    consumer
+        .subscribe(&[topic_name.as_str()])
+        .expect("could not subscribe to kafka topic");
 
-    let _consumer_future = consumer
+    consumer
         .stream()
-        .take(100)
+        .take(num_of_messages_to_send)
         .for_each(|message| {
             match message {
                 Ok(m) => {
                     let id = message_map[&(m.partition(), m.offset())];
                     match m.timestamp() {
                         Timestamp::CreateTime(timestamp) => assert!(timestamp >= start_time),
-                        _ => panic!("Expected createtime for message timestamp"),
+                        _ => panic!("Expected create time for message timestamp"),
                     };
-                    assert_eq!(m.payload_view::<str>().unwrap().unwrap(), value_fn(id));
-                    assert_eq!(m.key_view::<str>().unwrap().unwrap(), key_fn(id));
+                    assert_eq!(m.payload_view::<str>().unwrap().unwrap(), id.to_string());
+                    assert_eq!(m.key_view::<str>().unwrap().unwrap(), id.to_string());
                     assert_eq!(m.topic(), topic_name.as_str());
                 }
                 Err(e) => panic!("Error receiving message: {:?}", e),
@@ -107,13 +136,39 @@ async fn test_produce_consume_base() {
 /// waker slot.
 #[tokio::test(flavor = "multi_thread")]
 async fn test_produce_consume_base_concurrent() {
-    let _r = env_logger::try_init();
+    init_test_logger();
 
-    let topic_name = rand_test_topic("test_produce_consume_base_concurrent");
-    populate_topic(&topic_name, 100, &value_fn, &key_fn, None, None).await;
+    // Get Kafka container context.
+    let kafka_context = KafkaContext::shared()
+        .await
+        .expect("could not create kafka context");
 
-    let consumer = Arc::new(create_stream_consumer(&rand_test_group(), None));
-    consumer.subscribe(&[topic_name.as_str()]).unwrap();
+    let producer = producer::future_producer::create_producer(&kafka_context.bootstrap_servers)
+        .await
+        .expect("Could not create Future producer");
+
+    let num_of_messages_to_send = 100usize;
+    let topic_name = rand_test_topic("test_produce_consume_base_concurrent");
+    topics::populate_topic_using_future_producer(
+        &producer,
+        &topic_name,
+        num_of_messages_to_send,
+        None,
+    )
+        .await
+        .expect("Could not populate topic using Future producer");
+    // let message_map = populate_topic(&topic_name, 100, &value_fn, &key_fn, None, None).await;
+    let consumer = Arc::new(
+        consumer::stream_consumer::create_stream_consumer(
+            &kafka_context.bootstrap_servers,
+            Some(&rand_test_group()),
+        )
+        .await
+        .expect("could not create stream consumer"),
+    );
+    consumer
+        .subscribe(&[topic_name.as_str()])
+        .expect("could not subscribe to kafka topic");
 
     let mk_task = || {
         let consumer = consumer.clone();
@@ -137,13 +192,55 @@ async fn test_produce_consume_base_concurrent() {
 // All produced messages should be consumed.
 #[tokio::test(flavor = "multi_thread")]
 async fn test_produce_consume_base_assign() {
-    let _r = env_logger::try_init();
+    init_test_logger();
+
+    // Get Kafka container context.
+    let kafka_context = KafkaContext::shared()
+        .await
+        .expect("could not create kafka context");
 
     let topic_name = rand_test_topic("test_produce_consume_base_assign");
-    populate_topic(&topic_name, 10, &value_fn, &key_fn, Some(0), None).await;
-    populate_topic(&topic_name, 10, &value_fn, &key_fn, Some(1), None).await;
-    populate_topic(&topic_name, 10, &value_fn, &key_fn, Some(2), None).await;
-    let consumer = create_stream_consumer(&rand_test_group(), None);
+    let admin_client = admin::create_admin_client(&kafka_context.bootstrap_servers)
+        .await
+        .expect("Could not create admin client");
+    admin_client.create_topics(&new_topic_vec(&topic_name, Some(3)), &AdminOptions::default()).await.expect("could not create topics");
+
+    let producer = producer::future_producer::create_producer(&kafka_context.bootstrap_servers)
+        .await
+        .expect("Could not create Future producer");
+
+    let num_of_messages_to_send = 10usize;
+    topics::populate_topic_using_future_producer(
+        &producer,
+        &topic_name,
+        num_of_messages_to_send,
+        Some(0),
+    )
+        .await
+        .expect("Could not populate topic using Future producer");
+    topics::populate_topic_using_future_producer(
+        &producer,
+        &topic_name,
+        num_of_messages_to_send,
+        Some(1),
+    )
+        .await
+        .expect("Could not populate topic using Future producer");
+    topics::populate_topic_using_future_producer(
+        &producer,
+        &topic_name,
+        num_of_messages_to_send,
+        Some(2),
+    )
+        .await
+        .expect("Could not populate topic using Future producer");
+
+    let consumer = utils::consumer::stream_consumer::create_stream_consumer(
+        &kafka_context.bootstrap_servers,
+        Some(&rand_test_group()),
+    )
+    .await
+    .expect("could not create stream consumer");
     let mut tpl = TopicPartitionList::new();
     tpl.add_partition_offset(&topic_name, 0, Offset::Beginning)
         .unwrap();

tests/utils/admin.rs

@@ -17,7 +17,7 @@ pub async fn create_topic(
     topic_name: &'_ str,
 ) -> anyhow::Result<String> {
     let topic_results = admin_client
-        .create_topics(&new_topic_vec(&topic_name), &AdminOptions::default())
+        .create_topics(&new_topic_vec(&topic_name, None), &AdminOptions::default())
         .await
         .context("error creating topics")?;
     for topic_result in topic_results {
@@ -28,7 +28,7 @@ pub async fn create_topic(
     Ok(topic_name.to_string())
 }
 
-pub fn new_topic_vec(topic_name: &'_ str) -> Vec<NewTopic<'_>> {
-    let new_topic = NewTopic::new(&topic_name, 1, TopicReplication::Fixed(1));
+pub fn new_topic_vec(topic_name: &'_ str, num_partitions: Option<i32>) -> Vec<NewTopic<'_>> {
+    let new_topic = NewTopic::new(&topic_name, num_partitions.unwrap_or(1), TopicReplication::Fixed(1));
     vec![new_topic]
 }

tests/utils/consumer.rs renamed to tests/utils/consumer/mod.rs

@@ -1,3 +1,5 @@
+pub mod stream_consumer;
+
 use crate::utils::rand::rand_test_group;
 use anyhow::{bail, Context};
 use backon::{BlockingRetryable, ExponentialBuilder};

tests/utils/consumer/stream_consumer.rs

@@ -0,0 +1,19 @@
+use anyhow::Context;
+use rdkafka::ClientConfig;
+use rdkafka::config::FromClientConfig;
+use rdkafka::consumer::StreamConsumer;
+
+pub async fn create_stream_consumer(
+    bootstrap_server: &str,
+    consumer_group_option: Option<&str>,
+) -> anyhow::Result<StreamConsumer> {
+    let mut client_config = ClientConfig::default();
+    client_config.set("bootstrap.servers", bootstrap_server);
+    client_config.set("auto.offset.reset", "earliest");
+    if let Some(group) = consumer_group_option {
+        client_config.set("group.id", group);
+    }
+
+    let stream_consumer = StreamConsumer::from_config(&client_config).context("failed to create stream consumer")?;
+    Ok(stream_consumer)
+}

tests/utils/containers.rs

@@ -4,7 +4,7 @@ use std::sync::Arc;
 use testcontainers_modules::kafka::apache::Kafka;
 use testcontainers_modules::testcontainers::core::ContainerPort;
 use testcontainers_modules::testcontainers::runners::AsyncRunner;
-use testcontainers_modules::testcontainers::{ContainerAsync, Image, ImageExt};
+use testcontainers_modules::testcontainers::{ContainerAsync, Image};
 use tokio::sync::OnceCell;
 
 pub struct KafkaContext {

tests/utils/mod.rs

@@ -6,6 +6,7 @@ pub mod containers;
 pub mod logging;
 pub mod producer;
 pub mod rand;
+pub mod topics;
 
 use std::collections::HashMap;
 use std::env::{self};

tests/utils/topics.rs

@@ -0,0 +1,44 @@
+use rdkafka::producer::{FutureProducer, FutureRecord};
+use std::collections::HashMap;
+use std::time::Duration;
+
+pub type PartitionOffset = (i32, i64);
+pub type MessageId = usize;
+
+pub async fn populate_topic_using_future_producer(
+    producer: &FutureProducer,
+    topic_name: &str,
+    num_messages: usize,
+    partition: Option<i32>,
+) -> anyhow::Result<HashMap<PartitionOffset, MessageId>> {
+    let message_send_futures = (0..num_messages)
+        .map(|id| {
+            let future = async move {
+                producer
+                    .send(
+                        FutureRecord {
+                            topic: topic_name,
+                            payload: Some(&id.to_string()),
+                            key: Some(&id.to_string()),
+                            partition,
+                            timestamp: None,
+                            headers: None,
+                        },
+                        Duration::from_secs(1),
+                    )
+                    .await
+            };
+            (id, future)
+        })
+        .collect::<Vec<_>>();
+
+    let mut message_map = HashMap::<PartitionOffset, MessageId>::new();
+    for (id, future) in message_send_futures {
+        match future.await {
+            Ok(delivered) => message_map.insert((delivered.partition, delivered.offset), id),
+            Err((kafka_error, _message)) => panic!("Delivery failed: {}", kafka_error),
+        };
+    }
+
+    Ok(message_map)
+}