mod.rs

use std::future::Future;
use std::time::Duration;

use base64::{engine::general_purpose::URL_SAFE, Engine as _};
pub use fixtures::FixtureSnapshot;
use sha2::{Digest, Sha512};

use crate::connection::{ConnectOptions, Connection};
use crate::database::Database;
use crate::error::Error;
use crate::executor::Executor;
use crate::migrate::{Migrate, Migrator};
use crate::pool::{Pool, PoolConnection, PoolOptions};

mod fixtures;

pub trait TestSupport: Database {
    /// Get parameters to construct a `Pool` suitable for testing.
    ///
    /// This `Pool` instance will behave somewhat specially:
    /// * all handles share a single global semaphore to avoid exceeding the connection limit
    ///   on the database server.
    /// * each invocation results in a different temporary database.
    ///
    /// The implementation may require `DATABASE_URL` to be set in order to manage databases.
    /// The user credentials it contains must have the privilege to create and drop databases.
    fn test_context(
        args: &TestArgs,
    ) -> impl Future<Output = Result<TestContext<Self>, Error>> + Send + '_;

    fn cleanup_test(args: &TestArgs) -> impl Future<Output = Result<(), Error>> + Send + '_;

    /// Cleanup any test databases that are no longer in-use.
    ///
    /// Returns a count of the databases deleted, if possible.
    ///
    /// The implementation may require `DATABASE_URL` to be set in order to manage databases.
    /// The user credentials it contains must have the privilege to create and drop databases.
    fn cleanup_test_dbs() -> impl Future<Output = Result<Option<usize>, Error>> + Send + 'static;

    /// Cleanup any test databases that are no longer in-use.
    ///
    /// Returns a count of the databases deleted, if possible.
    fn cleanup_test_dbs_by_url(
        url: &str,
    ) -> impl Future<Output = Result<Option<usize>, Error>> + Send + '_;

    /// Take a snapshot of the current state of the database (data only).
    ///
    /// This snapshot can then be used to generate test fixtures.
    fn snapshot(
        conn: &mut Self::Connection,
    ) -> impl Future<Output = Result<FixtureSnapshot<Self>, Error>> + Send + '_;

    /// Generate a unique database name for the given test path.
    fn db_name(args: &TestArgs) -> String {
        let mut hasher = Sha512::new();
        hasher.update(args.test_path.as_bytes());
        let hash = hasher.finalize();
        let hash = URL_SAFE.encode(&hash[..39]);
        let db_name = format!("_sqlx_test_{}", hash).replace('-', "_");
        debug_assert!(db_name.len() == 63);
        db_name
    }
}

pub struct TestFixture {
    pub path: &'static str,
    pub contents: &'static str,
}

pub struct TestArgs {
    pub test_path: &'static str,
    pub migrator: Option<&'static Migrator>,
    pub fixtures: &'static [TestFixture],
    pub database_url_var: &'static str,
}

pub trait TestFn {
    type Output;

    fn run_test(self, args: TestArgs) -> Self::Output;
}

pub trait TestTermination {
    fn is_success(&self) -> bool;
}

pub struct TestContext<DB: Database> {
    pub pool_opts: PoolOptions<DB>,
    pub connect_opts: <DB::Connection as Connection>::Options,
    pub db_name: String,
}

impl<DB, Fut> TestFn for fn(Pool<DB>) -> Fut
where
    DB: TestSupport + Database,
    DB::Connection: Migrate,
    for<'c> &'c mut DB::Connection: Executor<'c, Database = DB>,
    Fut: Future,
    Fut::Output: TestTermination,
{
    type Output = Fut::Output;

    fn run_test(self, args: TestArgs) -> Self::Output {
        run_test_with_pool(args, self)
    }
}

impl<DB, Fut> TestFn for fn(PoolConnection<DB>) -> Fut
where
    DB: TestSupport + Database,
    DB::Connection: Migrate,
    for<'c> &'c mut DB::Connection: Executor<'c, Database = DB>,
    Fut: Future,
    Fut::Output: TestTermination,
{
    type Output = Fut::Output;

    fn run_test(self, args: TestArgs) -> Self::Output {
        run_test_with_pool(args, |pool| async move {
            let conn = pool
                .acquire()
                .await
                .expect("failed to acquire test pool connection");
            let res = (self)(conn).await;
            pool.close().await;
            res
        })
    }
}

impl<DB, Fut> TestFn for fn(PoolOptions<DB>, <DB::Connection as Connection>::Options) -> Fut
where
    DB: Database + TestSupport,
    DB::Connection: Migrate,
    for<'c> &'c mut DB::Connection: Executor<'c, Database = DB>,
    Fut: Future,
    Fut::Output: TestTermination,
{
    type Output = Fut::Output;

    fn run_test(self, args: TestArgs) -> Self::Output {
        run_test(args, self)
    }
}

impl<Fut> TestFn for fn() -> Fut
where
    Fut: Future,
{
    type Output = Fut::Output;

    fn run_test(self, args: TestArgs) -> Self::Output {
        assert!(
            args.fixtures.is_empty(),
            "fixtures cannot be applied for a bare function"
        );
        crate::rt::test_block_on(self())
    }
}

impl TestArgs {
    pub fn new(test_path: &'static str) -> Self {
        TestArgs {
            test_path,
            migrator: None,
            fixtures: &[],
            database_url_var: "DATABASE_URL",
        }
    }

    pub fn migrator(&mut self, migrator: &'static Migrator) {
        self.migrator = Some(migrator);
    }

    pub fn no_migrator(&mut self) {
        self.migrator = None;
    }

    pub fn fixtures(&mut self, fixtures: &'static [TestFixture]) {
        self.fixtures = fixtures;
    }

    pub fn database_url_var(&mut self, database_url_var: &'static str) {
        self.database_url_var = database_url_var;
    }
}

impl TestTermination for () {
    fn is_success(&self) -> bool {
        true
    }
}

impl<T, E> TestTermination for Result<T, E> {
    fn is_success(&self) -> bool {
        self.is_ok()
    }
}

fn run_test_with_pool<DB, F, Fut>(args: TestArgs, test_fn: F) -> Fut::Output
where
    DB: TestSupport,
    DB::Connection: Migrate,
    for<'c> &'c mut DB::Connection: Executor<'c, Database = DB>,
    F: FnOnce(Pool<DB>) -> Fut,
    Fut: Future,
    Fut::Output: TestTermination,
{
    let test_path = args.test_path;
    run_test::<DB, _, _>(args, |pool_opts, connect_opts| async move {
        let pool = pool_opts
            .connect_with(connect_opts)
            .await
            .expect("failed to connect test pool");

        let res = test_fn(pool.clone()).await;

        let close_timed_out = crate::rt::timeout(Duration::from_secs(10), pool.close())
            .await
            .is_err();

        if close_timed_out {
            eprintln!("test {test_path} held onto Pool after exiting");
        }

        res
    })
}

fn run_test<DB, F, Fut>(args: TestArgs, test_fn: F) -> Fut::Output
where
    DB: TestSupport,
    DB::Connection: Migrate,
    for<'c> &'c mut DB::Connection: Executor<'c, Database = DB>,
    F: FnOnce(PoolOptions<DB>, <DB::Connection as Connection>::Options) -> Fut,
    Fut: Future,
    Fut::Output: TestTermination,
{
    crate::rt::test_block_on(async move {
        let test_context = DB::test_context(&args)
            .await
            .expect("failed to connect to setup test database");

        setup_test_db::<DB>(&test_context.connect_opts, &args).await;

        let res = test_fn(test_context.pool_opts, test_context.connect_opts).await;

        if res.is_success() {
            if let Err(e) = DB::cleanup_test(&args).await {
                eprintln!(
                    "failed to delete database {:?}: {}",
                    test_context.db_name, e
                );
            }
        }

        res
    })
}

async fn setup_test_db<DB: Database>(
    copts: &<DB::Connection as Connection>::Options,
    args: &TestArgs,
) where
    DB::Connection: Migrate + Sized,
    for<'c> &'c mut DB::Connection: Executor<'c, Database = DB>,
{
    let mut conn = copts
        .connect()
        .await
        .expect("failed to connect to test database");

    if let Some(migrator) = args.migrator {
        migrator
            .run_direct(None, &mut conn)
            .await
            .expect("failed to apply migrations");
    }

    for fixture in args.fixtures {
        (&mut conn)
            .execute(fixture.contents)
            .await
            .unwrap_or_else(|e| panic!("failed to apply test fixture {:?}: {:?}", fixture.path, e));
    }

    conn.close()
        .await
        .expect("failed to close setup connection");
}

macro_rules! impl_test_fn {
    (
        $name:ident;
        $run_fn:ident;
        $run_with_pool_fn:ident;
        $(
            (
                $lt:lifetime $db:ident,
                $args:ident, $testctx:ident, $testpath:ident,
                $poolopts:ident, $connopts:ident,
                $pool:ident, $conn:ident
            ),
        )*;
    ) => {
            pub trait $name {
                type Output;

                fn run_test(self, $($args: TestArgs,)*) -> Self::Output;
            }

            impl<$($db,)* Fut> $name for fn($(Pool<$db>,)*) -> Fut
            where
                $(
                    $db: TestSupport + Database,
                    $db::Connection: Migrate,
                    for<$lt> &$lt mut $db::Connection: Executor<$lt, Database = $db>,
                )*
                Fut: Future,
                Fut::Output: TestTermination,
            {
                type Output = Fut::Output;

                fn run_test(self, $($args: TestArgs,)*) -> Self::Output {
                    $run_with_pool_fn($($args,)* self)
                }
            }

            impl<$($db,)* Fut> $name for fn($(PoolConnection<$db>,)*) -> Fut
            where
                $(
                    $db: TestSupport + Database,
                    $db::Connection: Migrate,
                    for<$lt> &$lt mut $db::Connection: Executor<$lt, Database = $db>,
                )*
                Fut: Future,
                Fut::Output: TestTermination,
            {
                type Output = Fut::Output;

                fn run_test(self, $($args: TestArgs,)*) -> Self::Output {
                    $run_with_pool_fn($($args,)* |$($pool,)*| async move {
                        $(
                            let $conn = $pool
                            .acquire()
                            .await
                            .expect("failed to acquire test pool connection");
                        )*


                        let res = (self)($($conn,)*).await;

                        $(
                            $pool.close().await;
                        )*

                        res
                    })
                }
            }

            impl<$($db,)* Fut> $name
                for fn(
                    $(
                        (PoolOptions<$db>, <$db::Connection as Connection>::Options),
                    )*
                ) -> Fut
            where
                $(
                    $db: TestSupport + Database,
                    $db::Connection: Migrate,
                    for<$lt> &$lt mut $db::Connection: Executor<$lt, Database = $db>,
                )*
                Fut: Future,
                Fut::Output: TestTermination,
            {
                type Output = Fut::Output;

                fn run_test(self, $($args: TestArgs,)*) -> Self::Output {
                    $run_fn($($args,)* self)
                }
            }

            impl<Fut> $name for fn() -> Fut
            where
                Fut: Future,
            {
                type Output = Fut::Output;

                fn run_test(self, $($args: TestArgs,)*) -> Self::Output {
                    $(
                        assert!(
                            $args.fixtures.is_empty(),
                            "fixtures cannot be applied for a bare function",
                        );
                    )*
                    crate::rt::test_block_on(self())
                }
            }

            fn $run_with_pool_fn<$($db,)* F, Fut>($($args: TestArgs,)* test_fn: F) -> Fut::Output
            where
                $(
                    $db: TestSupport,
                    $db::Connection: Migrate,
                    for<$lt> &$lt mut $db::Connection: Executor<$lt, Database = $db>,
                )*
                F: FnOnce($(Pool<$db>,)*) -> Fut,
                Fut: Future,
                Fut::Output: TestTermination,
            {
                $(
                    let $testpath: &'static str = $args.test_path;
                )*

                $run_fn::<$($db,)* _, _>(
                    $($args,)*
                    |$(($poolopts, $connopts),)*| async move {
                        $(
                            let $pool = $poolopts
                                .connect_with($connopts)
                                .await
                                .expect("failed to connect test pool");
                        )*

                        let res = test_fn($($pool.clone(),)*).await;

                        $(
                            let close_timed_out = crate::rt::timeout(Duration::from_secs(10), $pool.close())
                                .await
                                .is_err();
                            if close_timed_out {
                                eprintln!("test {} held onto Pool after exiting", $testpath);
                            }
                        )*

                        res
                    },
                )
            }

            fn $run_fn<$($db,)* F, Fut>($($args: TestArgs,)* test_fn: F) -> Fut::Output
            where
                $(
                    $db: TestSupport,
                    $db::Connection: Migrate,
                    for<$lt> &$lt mut $db::Connection: Executor<$lt, Database = $db>,
                )*
                F: FnOnce(
                    $((PoolOptions<$db>, <$db::Connection as Connection>::Options),)*
                ) -> Fut,
                Fut: Future,
                Fut::Output: TestTermination,
            {
                crate::rt::test_block_on(async move {
                    $(
                        let $testctx = $db::test_context(&$args)
                            .await
                            .expect("failed to connect to setup test database");
                        setup_test_db::<$db>(&$testctx.connect_opts, &$args).await;
                    )*

                    let res = test_fn(
                        $(($testctx.pool_opts, $testctx.connect_opts),)*
                    )
                    .await;

                    if res.is_success() {
                        $(
                            if let Err(e) = $db::cleanup_test(&$args).await {
                                eprintln!(
                                    "failed to delete database {:?}: {}",
                                    $testctx.db_name, e
                                );
                            }
                        )*
                    }

                    res
                })
            }

        };
}

impl_test_fn!(
    TestFn2;
    run_test2;
    run_test_with_pool2;
    ('c DB1, args1, tc1, tp1, po1, co1, p1, c1),
    ('d DB2, args2, tc2, tp2, po2, co2, p2, c2),
    ;
);

impl_test_fn!(
    TestFn3;
    run_test3;
    run_test_with_pool3;
    ('c DB1, args1, tc1, tp1, po1, co1, p1, c1),
    ('d DB2, args2, tc2, tp2, po2, co2, p2, c2),
    ('e DB3, args3, tc3, tp3, po3, co3, p3, c3),
    ;
);

impl_test_fn!(
    TestFn4;
    run_test4;
    run_test_with_pool4;
    ('c DB1, args1, tc1, tp1, po1, co1, p1, c1),
    ('d DB2, args2, tc2, tp2, po2, co2, p2, c2),
    ('e DB3, args3, tc3, tp3, po3, co3, p3, c3),
    ('f DB4, args4, tc4, tp4, po4, co4, p4, c4),
    ;
);