pgmq 0.34.0-alpha.4

//! Integration tests for the [`Queue`] trait and its implementations.
//!
//! Tests are written with the help of the custom [`pgmq_test_macro::queue_test`] macro -- the macro
//! generates individual tests for each type that implements the [`Queue`] trait.
//!
//! In order to prevent tests generating conflicting data in the DB, each test creates a temporary
//! test DB for itself. The temporary DBs are automatically removed if the tests pass, unless the
//! `PGMQ_KEEP_TEST_DB` env var is set to `true`.

#![cfg(feature = "queue-experimental")]

use initialization::ConnDetails;
use pgmq::queue::Queue;
use pgmq::Message;
use rand::RngExt;
use serde_derive::{Deserialize, Serialize};
use serde_json::json;

static QUEUE: &str = "queue";

mod initialization {
    use chrono::Utc;
    use sqlx::AssertSqlSafe;
    use std::env;
    use std::sync::OnceLock;
    use url::Url;

    static KEEP_TEST_DB: OnceLock<bool> = OnceLock::new();
    static DB_URL: OnceLock<Url> = OnceLock::new();
    static TIMESTAMP: OnceLock<String> = OnceLock::new();

    const MAX_DB_NAME_LENGTH: usize = 63;

    fn test_db_name(test_name: &str) -> String {
        let timestamp = TIMESTAMP.get_or_init(|| Utc::now().timestamp().to_string());

        let db_name = format!("{test_name}/{timestamp}");
        if db_name.len() > MAX_DB_NAME_LENGTH {
            panic!("Test DB name `{db_name}` is too long! Max length is {MAX_DB_NAME_LENGTH}");
        }
        db_name
    }

    fn db_url() -> &'static Url {
        DB_URL.get_or_init(|| {
            let url = env::var("DATABASE_URL").unwrap_or_else(|_| {
                "postgres://postgres:postgres@localhost:5432/postgres".to_owned()
            });
            Url::parse(&url).unwrap()
        })
    }

    #[derive(Debug, Clone)]
    pub struct ConnDetails {
        pub original: &'static Url,
        pub test_db_name: String,
        pub test_db_url: Url,
    }

    impl ConnDetails {
        pub fn new() -> Self {
            let original = db_url();
            let test_name = std::thread::current().name().unwrap().to_string();
            let test_db_name = test_db_name(&test_name);
            let mut test_db_url = original.clone();
            test_db_url.set_path(&test_db_name);
            Self {
                original,
                test_db_name,
                test_db_url,
            }
        }

        async fn original_conn(&self) -> sqlx::postgres::PgConnection {
            sqlx_conn(self.original).await
        }
    }

    pub async fn before(conn_details: &ConnDetails) {
        let create_db_statement = format!("CREATE DATABASE \"{}\"", conn_details.test_db_name);
        sqlx::query(AssertSqlSafe(create_db_statement))
            .execute(&mut conn_details.original_conn().await)
            .await
            .unwrap();

        install_pgmq(conn_details).await;
    }

    async fn install_pgmq(conn_details: &ConnDetails) {
        // Todo: It's a little awkward to create an instance of `PGMQueueExt` just to init/install pgmq.
        //  In a future change, we could expand the `Queue` trait to include the init/install methods,
        //  then we could replace this method with a call to the client implementation.
        let queue = pgmq::PGMQueueExt::new(conn_details.test_db_url.to_string(), 1)
            .await
            .unwrap();

        #[cfg(feature = "install-sql-embedded")]
        let result = queue.install_sql_from_embedded().await.map(|_| true);
        #[cfg(not(feature = "install-sql"))]
        let result = queue.init().await;

        result.expect("failed to init pgmq");
    }

    pub async fn after(conn_details: &ConnDetails) {
        let keep_db = *KEEP_TEST_DB.get_or_init(|| {
            env::var("PGMQ_KEEP_TEST_DB")
                .ok()
                .and_then(|x| x.parse::<bool>().ok())
                .unwrap_or(false)
        });
        if keep_db {
            return;
        }
        let drop_db_statement = format!(
            "DROP DATABASE IF EXISTS \"{}\" WITH (FORCE)",
            conn_details.test_db_name
        );
        sqlx::query(AssertSqlSafe(drop_db_statement))
            .execute(&mut conn_details.original_conn().await)
            .await
            .unwrap();
    }

    pub async fn sqlx_conn(url: &Url) -> sqlx::postgres::PgConnection {
        use sqlx::ConnectOptions;
        sqlx::postgres::PgConnectOptions::from_url(url)
            .unwrap()
            .connect()
            .await
            .unwrap()
    }

    pub async fn pgmq_ext(url: &Url) -> pgmq::PGMQueueExt {
        pgmq::PGMQueueExt::new(url.to_string(), 2).await.unwrap()
    }
}

#[derive(Debug, Serialize, Deserialize, PartialEq, Eq)]
struct TestMessage {
    a: String,
    b: i32,
}

impl TestMessage {
    fn new() -> Self {
        Self {
            a: std::thread::current().name().unwrap().to_string(),
            b: rand::rng().random_range(0..i32::MAX),
        }
    }
}

#[pgmq_test_macro::queue_test]
async fn read(conn_details: ConnDetails, queue: impl Queue) {
    queue.create(QUEUE).await.unwrap();
    let msg = TestMessage::new();
    let msg_id = queue.send(QUEUE, &msg, json!({}), 0).await.unwrap();

    // The first read should read the message
    let read_msg: Message<TestMessage> = queue
        .read(QUEUE, 10, 1)
        .await
        .unwrap()
        .into_iter()
        .next()
        .unwrap();
    assert_eq!(msg_id, read_msg.msg_id);
    assert_eq!(read_msg.message, msg);

    // A second read should return no messages
    let read_msg: Vec<Message<TestMessage>> = queue.read(QUEUE, 10, 1).await.unwrap();
    assert!(read_msg.is_empty());
}

#[pgmq_test_macro::queue_test]
async fn read_multiple(conn_details: ConnDetails, queue: impl Queue) {
    queue.create(QUEUE).await.unwrap();
    let msg_id1 = queue
        .send(QUEUE, TestMessage::new(), json!({}), 0)
        .await
        .unwrap();
    let msg_id2 = queue
        .send(QUEUE, TestMessage::new(), json!({}), 0)
        .await
        .unwrap();

    let messages: Vec<Message<TestMessage>> = queue.read(QUEUE, 10, 2).await.unwrap();
    assert_eq!(2, messages.len());
    assert!(messages.iter().any(|msg| msg.msg_id == msg_id1));
    assert!(messages.iter().any(|msg| msg.msg_id == msg_id2));
}