s2_testcontainers/
lib.rs

1#![doc = include_str!("../README.md")]
2#![warn(missing_docs)]
3
4use std::time::Duration;
5
6use s2_sdk::{
7    S2,
8    error::RequestError,
9    types::{AccountEndpoint, BasinEndpoint, S2Config, S2Endpoints, ValidationError},
10};
11use testcontainers::{
12    ContainerAsync, ContainerRequest, GenericImage, ImageExt, TestcontainersError,
13    core::IntoContainerPort, runners::AsyncRunner,
14};
15use tokio::time::{Instant, sleep, timeout};
16
17/// Image repository for the S2 Docker image.
18pub const IMAGE: &str = "ghcr.io/s2-streamstore/s2";
19/// Default S2 image tag.
20pub const DEFAULT_TAG: &str = env!("CARGO_PKG_VERSION");
21/// Port exposed by s2-lite.
22pub const PORT: u16 = 80;
23/// Default access token used by [`S2Lite::client`].
24pub const DEFAULT_ACCESS_TOKEN: &str = "ignored";
25
26const HEALTH_TIMEOUT: Duration = Duration::from_secs(30);
27const HEALTH_POLL_INTERVAL: Duration = Duration::from_millis(100);
28const HEALTH_REQUEST_TIMEOUT: Duration = Duration::from_secs(2);
29
30/// Result type for this crate.
31pub type Result<T> = std::result::Result<T, Error>;
32
33/// Errors from s2-testcontainers helpers.
34#[derive(Debug, thiserror::Error)]
35pub enum Error {
36    /// Error from Testcontainers.
37    #[error("testcontainers error: {0}")]
38    Testcontainers(#[from] TestcontainersError),
39    /// Request error from the S2 SDK.
40    #[error("s2 sdk request error: {0}")]
41    Request(#[from] RequestError),
42    /// S2 endpoint or resource name validation error.
43    #[error("validation error: {0}")]
44    Validation(#[from] ValidationError),
45    /// s2-lite did not become healthy before the startup timeout.
46    #[error("s2-lite did not become healthy at {endpoint}")]
47    NotHealthy {
48        /// Endpoint that did not become healthy.
49        endpoint: String,
50    },
51}
52
53/// Running s2-lite Testcontainers instance.
54#[derive(Debug)]
55pub struct S2Lite {
56    container: ContainerAsync<GenericImage>,
57    endpoint: String,
58    client: S2,
59}
60
61impl S2Lite {
62    /// Start s2-lite with the default image tag.
63    pub async fn start() -> Result<Self> {
64        Self::start_with(DEFAULT_TAG).await
65    }
66
67    /// Start s2-lite with a specific image tag.
68    pub async fn start_with(tag: impl Into<String>) -> Result<Self> {
69        let container = s2_lite_image_with_tag(tag).start().await?;
70        let host = container.get_host().await?;
71        let port = container.get_host_port_ipv4(PORT).await?;
72        let endpoint = format!("http://{host}:{port}");
73
74        wait_until_healthy(&endpoint).await?;
75
76        let client = S2::new(s2_config_for_endpoint(&endpoint, DEFAULT_ACCESS_TOKEN)?)?;
77
78        Ok(Self {
79            container,
80            endpoint,
81            client,
82        })
83    }
84
85    /// Return the mapped HTTP endpoint for this s2-lite instance.
86    pub fn endpoint(&self) -> &str {
87        &self.endpoint
88    }
89
90    /// Build an [`S2Config`] for this s2-lite instance with the provided access token.
91    pub fn config(&self, access_token: impl Into<String>) -> Result<S2Config> {
92        s2_config_for_endpoint(&self.endpoint, access_token)
93    }
94
95    /// Build an [`S2`] client for this s2-lite instance.
96    pub fn client(&self) -> Result<S2> {
97        Ok(self.client.clone())
98    }
99
100    /// Return the underlying Testcontainers container.
101    pub fn container(&self) -> &ContainerAsync<GenericImage> {
102        &self.container
103    }
104}
105
106/// Return the default S2 Docker [`GenericImage`].
107pub fn s2_image() -> GenericImage {
108    s2_image_with_tag(DEFAULT_TAG)
109}
110
111/// Return an S2 Docker [`GenericImage`] with a specific tag.
112pub fn s2_image_with_tag(tag: impl Into<String>) -> GenericImage {
113    GenericImage::new(IMAGE.to_string(), tag.into())
114}
115
116/// Return the default S2 Docker [`ContainerRequest`] configured to run `s2 lite`.
117pub fn s2_lite_image() -> ContainerRequest<GenericImage> {
118    s2_lite_image_with_tag(DEFAULT_TAG)
119}
120
121/// Return an S2 Docker [`ContainerRequest`] with a specific tag configured to run `s2 lite`.
122pub fn s2_lite_image_with_tag(tag: impl Into<String>) -> ContainerRequest<GenericImage> {
123    s2_image_with_tag(tag)
124        .with_exposed_port(PORT.tcp())
125        .with_cmd(["lite"])
126}
127
128/// Build an [`S2Config`] wired to use an endpoint for both account and basin APIs.
129pub fn s2_config_for_endpoint(
130    endpoint: impl AsRef<str>,
131    access_token: impl Into<String>,
132) -> Result<S2Config> {
133    let endpoint = endpoint.as_ref();
134    let endpoints = S2Endpoints::new(
135        AccountEndpoint::new(endpoint)?,
136        BasinEndpoint::new(endpoint)?,
137    )?;
138
139    Ok(S2Config::new(access_token).with_endpoints(endpoints))
140}
141
142async fn wait_until_healthy(endpoint: &str) -> Result<()> {
143    let client = reqwest::Client::new();
144    let health_url = format!("{endpoint}/health");
145    let deadline = Instant::now() + HEALTH_TIMEOUT;
146
147    loop {
148        let now = Instant::now();
149        if now >= deadline {
150            return Err(Error::NotHealthy {
151                endpoint: endpoint.to_string(),
152            });
153        }
154
155        let request_timeout = HEALTH_REQUEST_TIMEOUT.min(deadline - now);
156        if let Ok(Ok(response)) = timeout(request_timeout, client.get(&health_url).send()).await
157            && response.status().is_success()
158        {
159            return Ok(());
160        }
161
162        let now = Instant::now();
163        if now >= deadline {
164            return Err(Error::NotHealthy {
165                endpoint: endpoint.to_string(),
166            });
167        }
168
169        sleep(HEALTH_POLL_INTERVAL.min(deadline - now)).await;
170    }
171}
172
173#[cfg(test)]
174mod tests {
175    use s2_sdk::types::{BasinName, EnsureBasinInput, EnsureStreamInput, StreamName};
176    use testcontainers::Image;
177
178    use super::*;
179
180    #[test]
181    fn s2_image_defaults_to_versioned_docker_image() {
182        let image = s2_image_with_tag("test-tag");
183
184        assert_eq!(image.name(), IMAGE);
185        assert_eq!(image.tag(), "test-tag");
186        assert!(image.expose_ports().is_empty());
187    }
188
189    #[test]
190    fn s2_lite_image_defaults_to_lite_command() {
191        let request = s2_lite_image_with_tag("test-tag");
192
193        assert_eq!(request.image().name(), IMAGE);
194        assert_eq!(request.image().tag(), "test-tag");
195        assert_eq!(request.image().expose_ports(), &[PORT.tcp()]);
196        assert_eq!(request.cmd().collect::<Vec<_>>(), ["lite"]);
197    }
198
199    #[tokio::test]
200    async fn config_uses_same_endpoint_for_account_and_basin() {
201        let config = s2_config_for_endpoint("http://localhost:8080", "ignored").unwrap();
202
203        S2::new(config).unwrap();
204    }
205
206    #[tokio::test]
207    async fn starts_s2_lite_and_ensures_resources() {
208        let s2 = S2Lite::start().await.unwrap();
209
210        let client = s2.client().unwrap();
211        let basin_name = "test-basin".parse::<BasinName>().unwrap();
212        client
213            .ensure_basin(EnsureBasinInput::new(basin_name.clone()))
214            .await
215            .unwrap();
216
217        let basin = client.basin(basin_name.clone());
218        let stream_name = "test-stream".parse::<StreamName>().unwrap();
219        basin
220            .ensure_stream(EnsureStreamInput::new(stream_name.clone()))
221            .await
222            .unwrap();
223
224        assert_eq!(basin_name.as_ref(), "test-basin");
225        assert_eq!(stream_name.as_ref(), "test-stream");
226    }
227}
s2_testcontainers/lib.rs

s2_testcontainers/
lib.rs
