pgdrift_db/
test_utils.rs

1use sqlx::PgPool;
2use testcontainers::{
3    GenericImage, ImageExt,
4    core::{ContainerAsync, WaitFor},
5    runners::AsyncRunner,
6};
7
8/// Test database container and pool wrapper
9pub struct TestDb {
10    pub pool: PgPool,
11    database_url: String,
12    _container: ContainerAsync<GenericImage>,
13}
14
15impl TestDb {
16    /// Start a new postgres container and connection pool
17    pub async fn new() -> Result<Self, Box<dyn std::error::Error>> {
18        let image = GenericImage::new("postgres", "16")
19            .with_wait_for(WaitFor::message_on_stderr(
20                "database system is ready to accept connections",
21            ))
22            .with_env_var("POSTGRES_USER", "pgdrift")
23            .with_env_var("POSTGRES_PASSWORD", "pgdrift_test")
24            .with_env_var("POSTGRES_DB", "pgdrift_test");
25
26        let container = image.start().await?;
27        let port = container.get_host_port_ipv4(5432).await?;
28
29        let database_url = format!(
30            "postgres://pgdrift:pgdrift_test@127.0.0.1:{}/pgdrift_test",
31            port
32        );
33
34        // Wait for the connection
35        let mut attempts = 0;
36        let pool = loop {
37            match PgPool::connect(&database_url).await {
38                Ok(p) => break p,
39                Err(_) if attempts < 30 => {
40                    attempts += 1;
41                    tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
42                }
43                Err(e) => return Err(Box::new(e)),
44            }
45        };
46
47        Ok(TestDb {
48            pool,
49            database_url,
50            _container: container,
51        })
52    }
53
54    /// Get the database URL for this test database
55    pub fn database_url(&self) -> &str {
56        &self.database_url
57    }
58
59    /// Cleanup all test tables
60    pub async fn cleanup(&self) -> Result<(), sqlx::Error> {
61        crate::fixtures::cleanup(&self.pool).await
62    }
63}
pgdrift_db/test_utils.rs

pgdrift_db/
test_utils.rs
