Skip to main content

cruxmont_pg_test_macro/
lib.rs

1//! We can use the macro with the following:
2//! ```
3//! #[db_test]
4//! async fn test_create_username_conflict() {
5//!     // SQLX_POSTGRES_POOL is the test DB pool provided by the macro
6//!     let pool: &Pool<Postgres> = &*SQLX_POSTGRES_POOL;
7//!     run_migrations(pool).await;
8//!
9//!     let all_users = SqlxPostGresDescriptor::get_all_user_profiles(pool)
10//!         .await
11//!         .expect("get all users");
12//!
13//!     assert_eq!(0, all_users.len());
14//!
15//!     let new_user = gen_new_user_schema();
16//!
17//!     // TestDbHandle is provided by the test macro and yields the SQLX_POSTGRES_POOL
18//!     let outcome =
19//!         create_customer::<SqlxPostGresDescriptor, TestDbHandle, TestConfig>(Json(new_user))
20//!             .await;
21//!     assert_eq!(201, outcome.into_response().status());
22//!     let mut new_user = gen_new_user_schema();
23//!     new_user.email = "another_email".into();
24//!     insert_duplicate::<TestDbHandle>(new_user, "username already exists").await;
25//! }
26//! ```
27extern crate proc_macro;
28use proc_macro::TokenStream;
29use proc_macro2::Span;
30use quote::quote;
31use syn::{ItemFn, LitStr, parse_macro_input};
32use uuid::Uuid;
33
34#[proc_macro_attribute]
35pub fn pg_test(_attr: TokenStream, item: TokenStream) -> TokenStream {
36    let input_fn = parse_macro_input!(item as ItemFn);
37
38    // Get the function name
39    let func_name = &input_fn.sig.ident;
40    let stmts = &input_fn.block.stmts; // Vec<Stmt>
41
42    let db_tag = Uuid::new_v4().simple().to_string(); // "c0a801e6be7d470691ed7d087b4e1bfd"
43    let env_name = format!("DATABASE_URL_{}", db_tag); // "DATABASE_URL_c0a801e6be7d…"
44    let env_lit = LitStr::new(&env_name, Span::call_site()); // `syn::LitStr` → "DATABASE_URL_c0…"
45
46    let db_url = format!("postgres://username:password@localhost:5433/{db_tag}");
47    let db_name = LitStr::new(&db_tag, Span::call_site());
48    let db_url_lit = LitStr::new(&db_url, Span::call_site());
49    let pg_url_lit = LitStr::new(
50        "postgres://username:password@localhost:5433/main_db",
51        Span::call_site(),
52    );
53
54    let expanded = quote! {
55        #[test]
56        fn #func_name() {
57            use sqlx::Executor;
58            use tokio::runtime::Builder;
59
60            // set the environment variables for the DB
61            unsafe {
62                std::env::set_var(#env_lit, #db_url_lit);
63                std::env::set_var("DB_MAX_CONNECTIONS", "1");
64            }
65
66            let rt = Builder::new_current_thread()
67                .enable_all()
68                .build()
69                .expect("create Tokio runtime");
70
71            // create the DB in the core DB
72            rt.block_on(async {
73                let db_name = #db_name;
74                let mut admin = sqlx::postgres::PgPoolOptions::new()
75                .max_connections(1)
76                .connect(#pg_url_lit)
77                .await.expect("make admin DB connection pool");
78                admin.execute(format!(r#"CREATE DATABASE "{db_name}""#).as_str()).await.expect("execute create DB transaction");
79                admin.close().await;
80            });
81
82            let test_result = rt.block_on(async {
83                // define the DB pool which is std::sync::LazyLock<sqlx::Pool<sqlx::Postgres>> inner = Pool<sqlx::Postgres>
84                cruxmont::pg_pool::define_pg_pool!(SQLX_POSTGRES_TEST_POOL, #env_lit, "DB_MAX_CONNECTIONS");
85
86                struct TestDbHandle;
87
88                impl cruxmont::dal::connections::sqlx_postgres::YieldPostGresPool for TestDbHandle {
89                    fn yield_pool() -> &'static sqlx::Pool<sqlx::Postgres> {
90                        &*SQLX_POSTGRES_TEST_POOL
91                    }
92                }
93
94                // execute the testing code
95                let handle = tokio::spawn(async {
96                    #(#stmts)*
97                });
98                handle.await   // → Result<Result<(), _>, JoinError>
99            });
100
101            // drop the DB table
102            rt.block_on(async {
103                let db_name = #db_name;
104                let mut admin = sqlx::postgres::PgPoolOptions::new()
105                .max_connections(1)
106                .connect(#pg_url_lit)
107                .await.expect("make admin DB connection pool");
108                admin.execute(format!(r#"DROP DATABASE "{db_name}" WITH (FORCE)"#).as_str()).await.expect("execute drop DB transaction");
109                admin.close().await;
110            });
111            test_result.unwrap();
112        }
113    };
114    TokenStream::from(expanded)
115}