allow concurrency between tests

Author: xvello

Cargo.lock

@@ -5,6 +5,7 @@ members = [
 ]
 
 [workspace.dependencies]
+assert-json-diff = "2.0.2"
 axum = "0.6.15"
 axum-client-ip = "0.4.1"
 tokio = { version = "1.0", features = ["full"] }

Cargo.toml

@@ -13,3 +13,11 @@ time = { workspace = true }
 envconfig = "0.10.0"
 anyhow = { workspace = true, features = [] }
 once_cell = "1.18.0"
+
+[dev-dependencies]
+assert-json-diff =  { workspace = true }
+futures = "0.3.29"
+rand = { workspace = true }
+rdkafka = { workspace = true }
+reqwest = "0.11.22"
+serde_json = { workspace = true }

capture-server/Cargo.toml

@@ -8,6 +8,8 @@ pub struct Config {
     #[envconfig(default = "127.0.0.1:3000")]
     pub address: SocketAddr,
     pub redis_url: String,
+    #[envconfig(default = "true")]
+    pub export_prometheus: bool,
 
     pub kafka_hosts: String,
     pub kafka_topic: String,

capture-server/src/config.rs

@@ -23,7 +23,7 @@ where
             sink::PrintSink {},
             redis_client,
             billing,
-            true,
+            config.export_prometheus,
         )
     } else {
         let sink =
@@ -34,7 +34,7 @@ where
             sink,
             redis_client,
             billing,
-            true,
+            config.export_prometheus,
         )
     };
 

capture-server/src/server.rs

@@ -0,0 +1,169 @@
+#![allow(dead_code)]
+
+use std::default::Default;
+use std::net::{SocketAddr, TcpListener};
+use std::str::FromStr;
+use std::string::ToString;
+use std::sync::{Arc, Once};
+use std::time::Duration;
+
+use anyhow::bail;
+use once_cell::sync::Lazy;
+use rand::distributions::Alphanumeric;
+use rand::Rng;
+use rdkafka::admin::{AdminClient, AdminOptions, NewTopic, TopicReplication};
+use rdkafka::config::{ClientConfig, FromClientConfig};
+use rdkafka::consumer::{BaseConsumer, Consumer};
+use rdkafka::util::Timeout;
+use rdkafka::{Message, TopicPartitionList};
+use tokio::sync::Notify;
+use tracing::debug;
+
+use capture_server::config::Config;
+use capture_server::server::serve;
+
+pub static DEFAULT_CONFIG: Lazy<Config> = Lazy::new(|| Config {
+    print_sink: false,
+    address: SocketAddr::from_str("127.0.0.1:0").unwrap(),
+    export_prometheus: false,
+    redis_url: "redis://localhost:6379/".to_string(),
+    kafka_hosts: "kafka:9092".to_string(),
+    kafka_topic: "events_plugin_ingestion".to_string(),
+    kafka_tls: false,
+});
+
+static TRACING_INIT: Once = Once::new();
+pub fn setup_tracing() {
+    TRACING_INIT.call_once(|| {
+        tracing_subscriber::fmt()
+            .with_writer(tracing_subscriber::fmt::TestWriter::new())
+            .init()
+    });
+}
+pub struct ServerHandle {
+    pub addr: SocketAddr,
+    shutdown: Arc<Notify>,
+}
+
+impl ServerHandle {
+    pub fn for_topic(topic: &EphemeralTopic) -> Self {
+        let mut config = DEFAULT_CONFIG.clone();
+        config.kafka_topic = topic.topic_name().to_string();
+        Self::for_config(config)
+    }
+    pub fn for_config(config: Config) -> Self {
+        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
+        let addr = listener.local_addr().unwrap();
+        let notify = Arc::new(Notify::new());
+        let shutdown = notify.clone();
+
+        tokio::spawn(
+            async move { serve(config, listener, async { notify.notified().await }).await },
+        );
+        Self { addr, shutdown }
+    }
+
+    pub async fn capture_events<T: Into<reqwest::Body>>(&self, body: T) -> reqwest::Response {
+        let client = reqwest::Client::new();
+        client
+            .post(format!("http://{:?}/i/v0/e", self.addr))
+            .body(body)
+            .send()
+            .await
+            .expect("failed to send request")
+    }
+}
+
+impl Drop for ServerHandle {
+    fn drop(&mut self) {
+        self.shutdown.notify_one()
+    }
+}
+
+pub struct EphemeralTopic {
+    consumer: BaseConsumer,
+    read_timeout: Timeout,
+    topic_name: String,
+}
+
+impl EphemeralTopic {
+    pub async fn new() -> Self {
+        let mut config = ClientConfig::new();
+        config.set("group.id", "capture_integration_tests");
+        config.set("bootstrap.servers", DEFAULT_CONFIG.kafka_hosts.clone());
+        config.set("debug", "all");
+
+        // TODO: check for name collision?
+        let topic_name = random_string("events_", 16);
+        let admin = AdminClient::from_config(&config).expect("failed to create admin client");
+        admin
+            .create_topics(
+                &[NewTopic {
+                    name: &topic_name,
+                    num_partitions: 1,
+                    replication: TopicReplication::Fixed(1),
+                    config: vec![],
+                }],
+                &AdminOptions::default(),
+            )
+            .await
+            .expect("failed to create topic");
+
+        let consumer: BaseConsumer = config.create().expect("failed to create consumer");
+        let mut assignment = TopicPartitionList::new();
+        assignment.add_partition(&topic_name, 0);
+        consumer
+            .assign(&assignment)
+            .expect("failed to assign topic");
+
+        Self {
+            consumer,
+            read_timeout: Timeout::After(Duration::from_secs(5)),
+            topic_name,
+        }
+    }
+
+    pub fn next_event(&self) -> anyhow::Result<serde_json::Value> {
+        match self.consumer.poll(self.read_timeout) {
+            Some(Ok(message)) => {
+                let body = message.payload().expect("empty kafka message");
+                let event = serde_json::from_slice(body)?;
+                Ok(event)
+            }
+            Some(Err(err)) => bail!("kafka read error: {}", err),
+            None => bail!("kafka read timeout"),
+        }
+    }
+
+    pub fn topic_name(&self) -> &str {
+        &self.topic_name
+    }
+}
+
+impl Drop for EphemeralTopic {
+    fn drop(&mut self) {
+        debug!("dropping EphemeralTopic {}...", self.topic_name);
+        _ = self.consumer.unassign();
+        futures::executor::block_on(delete_topic(self.topic_name.clone()));
+        debug!("dropped topic");
+    }
+}
+
+async fn delete_topic(topic: String) {
+    let mut config = ClientConfig::new();
+    config.set("bootstrap.servers", DEFAULT_CONFIG.kafka_hosts.clone());
+    let admin = AdminClient::from_config(&config).expect("failed to create admin client");
+    admin
+        .delete_topics(&[&topic], &AdminOptions::default())
+        .await
+        .expect("failed to delete topic");
+}
+
+pub fn random_string(prefix: &str, length: usize) -> String {
+    let suffix: String = rand::thread_rng()
+        .sample_iter(Alphanumeric)
+        .take(length)
+        .map(char::from)
+        .collect();
+    format!("{}_{}", prefix, suffix)
+}

capture-server/tests/common.rs

@@ -1,54 +1,75 @@
-use std::net::{SocketAddr, TcpListener};
-use std::str::FromStr;
-use std::string::ToString;
-
 use anyhow::Result;
-use once_cell::sync::Lazy;
-use tokio::sync::oneshot;
-use tokio::task::JoinHandle;
-
-use capture_server::config::Config;
-use capture_server::server::serve;
-
-static DEFAULT_CONFIG: Lazy<Config> = Lazy::new(|| Config {
-    print_sink: false,
-    address: SocketAddr::from_str("127.0.0.1:0").unwrap(),
-    redis_url: "redis://localhost:6379/".to_string(),
-    kafka_hosts: "kafka:9092".to_string(),
-    kafka_topic: "events_plugin_ingestion".to_string(),
-    kafka_tls: false,
-});
-
-struct ServerHandle {
-    addr: SocketAddr,
-    shutdown: oneshot::Sender<()>,
-    join: JoinHandle<()>,
-}
+use assert_json_diff::assert_json_include;
+use reqwest::StatusCode;
+use serde_json::json;
 
-impl ServerHandle {
-    fn new(config: Config) -> Self {
-        let (shutdown, rx) = oneshot::channel::<()>();
-        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
-        let addr = listener.local_addr().unwrap();
-        let join =
-            tokio::spawn(async move { serve(config, listener, async { rx.await.unwrap() }).await });
-        Self {
-            addr,
-            shutdown,
-            join,
-        }
-    }
-
-    async fn stop(self) -> Result<()> {
-        self.shutdown.send(()).unwrap();
-        self.join.await?;
-        Ok(())
-    }
-}
+use crate::common::*;
+mod common;
 
 #[tokio::test]
 async fn it_captures_one_event() -> Result<()> {
-    let server = ServerHandle::new(DEFAULT_CONFIG.clone());
+    setup_tracing();
+    let token = random_string("token", 16);
+    let distinct_id = random_string("id", 16);
+    let topic = EphemeralTopic::new().await;
+    let server = ServerHandle::for_topic(&topic);
+
+    let event = json!({
+        "token": token,
+        "event": "testing",
+        "distinct_id": distinct_id
+    });
+    let res = server.capture_events(event.to_string()).await;
+    assert_eq!(StatusCode::OK, res.status());
+
+    let event = topic.next_event()?;
+    assert_json_include!(
+        actual: event,
+        expected: json!({
+            "token": token,
+            "distinct_id": distinct_id
+        })
+    );
+
+    Ok(())
+}
+
+#[tokio::test]
+async fn it_captures_a_batch() -> Result<()> {
+    setup_tracing();
+    let token = random_string("token", 16);
+    let distinct_id1 = random_string("id", 16);
+    let distinct_id2 = random_string("id", 16);
+
+    let topic = EphemeralTopic::new().await;
+    let server = ServerHandle::for_topic(&topic);
+
+    let event = json!([{
+        "token": token,
+        "event": "event1",
+        "distinct_id": distinct_id1
+    },{
+        "token": token,
+        "event": "event2",
+        "distinct_id": distinct_id2
+    }]);
+    let res = server.capture_events(event.to_string()).await;
+    assert_eq!(StatusCode::OK, res.status());
+
+    assert_json_include!(
+        actual: topic.next_event()?,
+        expected: json!({
+            "token": token,
+            "distinct_id": distinct_id1
+        })
+    );
+    assert_json_include!(
+        actual: topic.next_event()?,
+        expected: json!({
+            "token": token,
+            "distinct_id": distinct_id2
+        })
+    );
 
-    server.stop().await
+    Ok(())
 }

capture-server/tests/events.rs

@@ -32,7 +32,7 @@ thiserror = { workspace = true }
 redis = { version="0.23.3", features=["tokio-comp", "cluster", "cluster-async"] }
 
 [dev-dependencies]
-assert-json-diff = "2.0.2"
+assert-json-diff =  { workspace = true }
 axum-test-helper = "0.2.0"
 mockall = "0.11.2"
 redis-test = "0.2.3"

capture/Cargo.toml

@@ -66,7 +66,6 @@ pub fn router<
     // does not work well.
     if metrics {
         let recorder_handle = setup_metrics_recorder();
-
         router.route("/metrics", get(move || ready(recorder_handle.render())))
     } else {
         router

capture/src/router.rs

@@ -1,3 +1,5 @@
+version: "3"
+
 services:
     zookeeper:
         image: zookeeper:3.7.0
@@ -17,10 +19,33 @@ services:
             ALLOW_PLAINTEXT_LISTENER: 'true'
         ports:
             - '9092:9092'
+        healthcheck:
+            test: kafka-cluster.sh cluster-id --bootstrap-server localhost:9092 || exit 1
+            interval: 3s
+            timeout: 10s
+            retries: 10
 
     redis:
         image: redis:6.2.7-alpine
         restart: on-failure
         command: redis-server --maxmemory-policy allkeys-lru --maxmemory 200mb
         ports:
             - '6379:6379'
+        healthcheck:
+            test: ["CMD", "redis-cli", "ping"]
+            interval: 3s
+            timeout: 10s
+            retries: 10
+
+    kafka-ui:
+        image: provectuslabs/kafka-ui:latest
+        profiles: ["ui"]
+        ports:
+            - '8080:8080'
+        depends_on:
+            - zookeeper
+            - kafka
+        environment:
+            KAFKA_CLUSTERS_0_NAME: local
+            KAFKA_CLUSTERS_0_BOOTSTRAPSERVERS: kafka:9092
+            KAFKA_CLUSTERS_0_ZOOKEEPER: zookeeper:2181