Skip to main content

aster_forge_test/
postgres.rs

1//! Shared reusable PostgreSQL container for integration tests.
2//!
3//! The container provides isolated databases with automatic stale-resource cleanup. Products own
4//! their migrations and seed data; this module owns database creation, connection retry, and
5//! teardown mechanics.
6
7use crate::state::{ContainerLease, ContainerStateLock};
8use crate::suite::TestContainerSuite;
9use sea_orm::{ConnectionTrait, Database, DatabaseConnection};
10use std::time::Duration;
11use testcontainers::core::{ContainerAsync, IntoContainerPort, WaitFor};
12use testcontainers::{GenericImage, ImageExt, ReuseDirective, runners::AsyncRunner};
13
14/// Handle to the suite's shared PostgreSQL container.
15pub struct PostgresTestContainer {
16    admin_url: String,
17    suite: TestContainerSuite,
18    _container: ContainerAsync<GenericImage>,
19    _lease: ContainerLease,
20}
21
22/// Isolated PostgreSQL database owned by one test process.
23pub struct PostgresTestDatabase {
24    name: String,
25    url: String,
26    admin_url: String,
27    suite: TestContainerSuite,
28}
29
30impl PostgresTestContainer {
31    /// Starts (or reuses) the shared PostgreSQL container with `postgres`/`postgres` credentials.
32    pub async fn start(suite: &TestContainerSuite) -> Self {
33        let lock = ContainerStateLock::acquire(suite, "postgres");
34        let mut state = lock.load();
35        let stale_resources = state.prune_stale();
36        state.register_pid(std::process::id());
37        for resource in &stale_resources {
38            state.remember_resource(std::process::id(), resource);
39        }
40        lock.save(&state);
41
42        let container = GenericImage::new("postgres", "16")
43            .with_exposed_port(IntoContainerPort::tcp(5432))
44            .with_wait_for(WaitFor::message_on_stderr(
45                "database system is ready to accept connections",
46            ))
47            .with_container_name(suite.container_name("postgres"))
48            .with_reuse(ReuseDirective::Always)
49            .with_env_var("POSTGRES_USER", "postgres")
50            .with_env_var("POSTGRES_PASSWORD", "postgres")
51            .with_env_var("POSTGRES_DB", "postgres")
52            .start()
53            .await
54            .expect("failed to start PostgreSQL test container");
55        let port = container
56            .get_host_port_ipv4(IntoContainerPort::tcp(5432))
57            .await
58            .expect("PostgreSQL test port should be exposed");
59        drop(lock);
60
61        let fixture = Self {
62            admin_url: format!("postgres://postgres:postgres@127.0.0.1:{port}/postgres"),
63            suite: suite.clone(),
64            _container: container,
65            _lease: ContainerLease::new(suite.clone(), "postgres"),
66        };
67        fixture.cleanup_databases(&stale_resources).await;
68        fixture
69    }
70
71    /// Returns the admin URL pointing at the default `postgres` database.
72    pub fn admin_url(&self) -> &str {
73        &self.admin_url
74    }
75
76    /// Creates and registers an isolated database for a product test.
77    pub async fn create_database(&self, name: &str) -> PostgresTestDatabase {
78        assert_valid_database_name(name);
79        let lock = ContainerStateLock::acquire(&self.suite, "postgres");
80        let mut state = lock.load();
81        state.remember_resource(std::process::id(), name);
82        lock.save(&state);
83        drop(lock);
84
85        let admin = connect_with_retry(&self.admin_url).await;
86        admin
87            .execute_unprepared(&format!("CREATE DATABASE {}", quote_identifier(name)))
88            .await
89            .unwrap_or_else(|error| {
90                panic!("failed to create PostgreSQL test database {name}: {error}")
91            });
92        admin
93            .close()
94            .await
95            .unwrap_or_else(|error| panic!("failed to close PostgreSQL admin connection: {error}"));
96
97        PostgresTestDatabase {
98            name: name.to_string(),
99            url: database_url(&self.admin_url, name),
100            admin_url: self.admin_url.clone(),
101            suite: self.suite.clone(),
102        }
103    }
104
105    async fn cleanup_databases(&self, names: &[String]) {
106        if names.is_empty() {
107            return;
108        }
109        let admin = connect_with_retry(&self.admin_url).await;
110        for name in names {
111            admin
112                .execute_unprepared(&format!(
113                    "DROP DATABASE IF EXISTS {} WITH (FORCE)",
114                    quote_identifier(name)
115                ))
116                .await
117                .unwrap_or_else(|error| {
118                    panic!("failed to drop stale PostgreSQL test database {name}: {error}")
119                });
120            let lock = ContainerStateLock::acquire(&self.suite, "postgres");
121            let mut state = lock.load();
122            state.forget_resource(std::process::id(), name);
123            lock.save(&state);
124        }
125        admin
126            .close()
127            .await
128            .unwrap_or_else(|error| panic!("failed to close PostgreSQL admin connection: {error}"));
129    }
130}
131
132impl PostgresTestDatabase {
133    /// Returns the isolated database name.
134    pub fn name(&self) -> &str {
135        &self.name
136    }
137
138    /// Returns the connection URL for this database.
139    pub fn url(&self) -> &str {
140        &self.url
141    }
142
143    /// Connects to this database, retrying while the service becomes ready.
144    pub async fn connect(&self) -> DatabaseConnection {
145        connect_with_retry(&self.url).await
146    }
147
148    /// Drops this database and removes it from the shared resource registry.
149    pub async fn cleanup(&self) {
150        let admin = connect_with_retry(&self.admin_url).await;
151        admin
152            .execute_unprepared(&format!(
153                "DROP DATABASE IF EXISTS {} WITH (FORCE)",
154                quote_identifier(&self.name)
155            ))
156            .await
157            .unwrap_or_else(|error| {
158                panic!(
159                    "failed to drop PostgreSQL test database {}: {error}",
160                    self.name
161                )
162            });
163        admin
164            .close()
165            .await
166            .unwrap_or_else(|error| panic!("failed to close PostgreSQL admin connection: {error}"));
167
168        let lock = ContainerStateLock::acquire(&self.suite, "postgres");
169        let mut state = lock.load();
170        state.forget_resource(std::process::id(), &self.name);
171        lock.save(&state);
172    }
173}
174
175async fn connect_with_retry(database_url: &str) -> DatabaseConnection {
176    let deadline = tokio::time::Instant::now() + Duration::from_secs(60);
177    loop {
178        match Database::connect(database_url).await {
179            Ok(database) => return database,
180            Err(error) if tokio::time::Instant::now() >= deadline => {
181                panic!("PostgreSQL test database did not become ready: {error}")
182            }
183            Err(_) => tokio::time::sleep(Duration::from_millis(250)).await,
184        }
185    }
186}
187
188fn database_url(admin_url: &str, name: &str) -> String {
189    admin_url
190        .rsplit_once('/')
191        .map(|(base, _)| format!("{base}/{name}"))
192        .unwrap_or_else(|| admin_url.to_string())
193}
194
195fn quote_identifier(value: &str) -> String {
196    format!("\"{}\"", value.replace('"', "\"\""))
197}
198
199fn assert_valid_database_name(name: &str) {
200    assert!(
201        !name.is_empty()
202            && name.len() <= 63
203            && name
204                .bytes()
205                .all(|byte| byte.is_ascii_alphanumeric() || byte == b'_'),
206        "PostgreSQL test database name must be 1-63 ASCII alphanumeric or '_' characters: {name:?}"
207    );
208}
209
210#[cfg(test)]
211mod tests {
212    use super::{assert_valid_database_name, database_url, quote_identifier};
213
214    #[test]
215    fn database_url_replaces_admin_database() {
216        assert_eq!(
217            database_url("postgres://user:pass@127.0.0.1:5432/postgres", "isolated"),
218            "postgres://user:pass@127.0.0.1:5432/isolated"
219        );
220    }
221
222    #[test]
223    fn identifier_quoting_escapes_quotes() {
224        assert_eq!(quote_identifier("test\"name"), "\"test\"\"name\"");
225    }
226
227    #[test]
228    fn database_name_accepts_boundaries() {
229        assert_valid_database_name("a");
230        assert_valid_database_name(&"a".repeat(63));
231        assert_valid_database_name("aster_product_123");
232    }
233
234    #[test]
235    fn database_name_rejects_unsafe_or_oversized_values() {
236        for name in ["", "has-hyphen", "has quote\"", "has space"] {
237            assert!(
238                std::panic::catch_unwind(|| assert_valid_database_name(name)).is_err(),
239                "database name {name:?} should be rejected"
240            );
241        }
242        let oversized = "a".repeat(64);
243        assert!(std::panic::catch_unwind(|| assert_valid_database_name(&oversized)).is_err());
244    }
245}