datum-mq 0.11.2

use std::{
    collections::BTreeSet,
    env,
    path::PathBuf,
    process::Command,
    time::{Duration, SystemTime, UNIX_EPOCH},
};

use bytes::Bytes;
use datum::{Sink, Source, StreamError};
use datum_mq::{
    ConsumerRecord, KafkaConsumerBackend, KafkaConsumerSettings, KafkaOffset, KafkaProducerBackend,
    KafkaProducerSettings, KafkaSink, KafkaSource, ProducerRecord, Subscription,
};

fn kafka_enabled() -> bool {
    env::var("DATUM_MQ_TEST_KAFKA").ok().as_deref() == Some("1")
}

fn security_enabled() -> bool {
    env::var("DATUM_MQ_TEST_KAFKA_SECURITY").ok().as_deref() == Some("1")
}

fn unique_topic(prefix: &str) -> String {
    let nanos = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .expect("system time after unix epoch")
        .as_nanos();
    format!("datum-mq-featureless-{prefix}-{nanos}")
}

fn mq_root() -> PathBuf {
    PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("../../baselines/mq")
}

fn kafka_topics() -> PathBuf {
    mq_root().join(".runtime/kafka_2.13-4.2.0/bin/kafka-topics.sh")
}

fn create_topic(bootstrap: &str, topic: &str, command_config: Option<PathBuf>) {
    let mut command = Command::new(kafka_topics());
    command
        .arg("--bootstrap-server")
        .arg(bootstrap)
        .arg("--create")
        .arg("--if-not-exists")
        .arg("--topic")
        .arg(topic)
        .arg("--partitions")
        .arg("2")
        .arg("--replication-factor")
        .arg("1");
    if let Some(command_config) = command_config {
        command.arg("--command-config").arg(command_config);
    }
    assert!(command.status().expect("run kafka-topics.sh").success());
}

fn produce_and_consume(
    topic: &str,
    producer: KafkaProducerSettings,
    consumer: KafkaConsumerSettings,
) {
    const COUNT: usize = 32;
    let records = (0..COUNT).map(|index| {
        ProducerRecord::new(topic, Some(Bytes::from(format!("value-{index}"))))
            .with_key(Some(Bytes::from(index.to_string())))
    });
    assert_eq!(producer.producer_backend, KafkaProducerBackend::Native);
    assert_eq!(producer.config.get("enable.idempotence"), Some("true"));
    let control = Source::from_iterable(records)
        .run_with(KafkaSink::plain(producer))
        .expect("native producer materializes");
    control
        .drain_and_shutdown()
        .expect("native producer drains");

    assert_eq!(consumer.consumer_backend, KafkaConsumerBackend::Native);
    let values = KafkaSource::committable(consumer, Subscription::topics([topic]))
        .as_source()
        .take(COUNT)
        .run_with(Sink::fold_result(
            BTreeSet::new(),
            |mut values, (record, offset): (ConsumerRecord, KafkaOffset)| {
                values.insert(record.payload.unwrap_or_default());
                offset.commit().map_err(StreamError::from)?;
                Ok(values)
            },
        ))
        .expect("native consumer materializes")
        .wait()
        .expect("native consumer completes");
    assert_eq!(values.len(), COUNT);
}

#[test]
fn native_paths_round_trip() {
    if !kafka_enabled() {
        eprintln!("skipping native_paths_round_trip; start baselines/mq/up.sh");
        return;
    }
    let bootstrap =
        env::var("MQ_BOOTSTRAP_SERVERS").unwrap_or_else(|_| "127.0.0.1:9092".to_owned());
    let topic = unique_topic("plaintext");
    create_topic(&bootstrap, &topic, None);
    produce_and_consume(
        &topic,
        KafkaProducerSettings::new(bootstrap.clone())
            .with("linger.ms", "0")
            .with_drain_timeout(Duration::from_secs(20)),
        KafkaConsumerSettings::new(bootstrap, format!("{topic}-group"))
            .with("auto.offset.reset", "earliest")
            .with_commit_batch_size(1)
            .with_drain_timeout(Duration::from_secs(20)),
    );
}

#[test]
fn secured_native_paths_round_trip() {
    if !security_enabled() {
        eprintln!("skipping secured_native_paths_round_trip; start baselines/mq/up-secure.sh");
        return;
    }
    let bootstrap =
        env::var("MQ_SASL_SSL_BOOTSTRAP_SERVERS").unwrap_or_else(|_| "localhost:9095".to_owned());
    let topic = unique_topic("sasl-ssl");
    create_topic(
        &bootstrap,
        &topic,
        Some(mq_root().join(".runtime/secure/sasl-plain-client.properties")),
    );
    let ca = mq_root().join(".runtime/secure/certs/ca.pem");
    let secure = |settings: KafkaProducerSettings| {
        settings
            .with("security.protocol", "sasl_ssl")
            .with("ssl.ca.location", ca.to_string_lossy())
            .with("sasl.mechanism", "PLAIN")
            .with("sasl.username", "plain")
            .with("sasl.password", "plain-secret")
    };
    let producer = secure(
        KafkaProducerSettings::new(bootstrap.clone())
            .with("linger.ms", "0")
            .with_drain_timeout(Duration::from_secs(20)),
    );
    let consumer = KafkaConsumerSettings::new(bootstrap, format!("{topic}-group"))
        .with("security.protocol", "sasl_ssl")
        .with("ssl.ca.location", ca.to_string_lossy())
        .with("sasl.mechanism", "PLAIN")
        .with("sasl.username", "plain")
        .with("sasl.password", "plain-secret")
        .with("auto.offset.reset", "earliest")
        .with_commit_batch_size(1)
        .with_drain_timeout(Duration::from_secs(20));
    produce_and_consume(&topic, producer, consumer);
}