Skip to main content

omry_integration/
flora_image.rs

1use secrecy::{ExposeSecret, SecretString};
2use std::borrow::Cow;
3use std::path::Path;
4use testcontainers::core::{ContainerPort, ContainerRequest, IntoContainerPort, Mount, WaitFor};
5use testcontainers::runners::AsyncRunner;
6use testcontainers::{ContainerAsync, Image, ImageExt};
7use tracing::Level;
8use url::Url;
9
10use crate::containers::{ContainerConfig, ContainerLogConsumer, ImageConfig};
11use crate::env_keys;
12use crate::helpers::{ContainerMount, HostPorts, TempPathWithString};
13
14/// The path to the data directory used by Omry inside the container.
15const DATA_PATH_CONTAINER: &str = "/usr/local/etc/flora";
16
17/// This image name assumes the image is stored in the local registry.
18const IMAGE_NAME: &str = "localhost:8777/omry-server";
19
20/// Omry server image tag.
21const IMAGE_TAG: &str = "latest";
22
23/// The port the Omry server should bind to inside the container for tarpc.
24const CONTAINER_PORT_RPC: u16 = 47473;
25
26/// The web/HTTP port the Omry server should bind to inside the container.  
27const CONTAINER_PORT_WEB: u16 = 47474;
28
29/// The message the Omry server should print to stdout when it's ready to receive requests.
30const READY_MESSAGE: &str = "flora is ready!";
31
32/// Enum to encode plain or secret environment variables
33/// for the flora server image.
34#[derive(Debug)]
35pub enum EnvVar {
36    /// Plain (not secret) environment variable.
37    Plain {
38        /// Key (variable name).
39        key: String,
40
41        /// Value.
42        value: String,
43    },
44
45    /// Secret environment variable, whose value should not be unintentionally exposed.
46    Secret {
47        /// Key (variable name).
48        key: String,
49
50        /// Value (will be stored in [`SecretString`]).
51        value: SecretString,
52    },
53}
54
55impl EnvVar {
56    /// Returns this environment variable as a string (key, value) pair.
57    fn as_key_value_pair(&self) -> (&str, &str) {
58        match self {
59            Self::Plain { key, value } => (key, value),
60            Self::Secret { key, value } => (key, value.expose_secret()),
61        }
62    }
63}
64
65/// flora image config.
66#[allow(clippy::missing_docs_in_private_items)]
67#[derive(Debug)]
68pub struct FloraConfig {
69    auth_token: EnvVar,
70    database_url: EnvVar,
71    ready_message: EnvVar,
72    rpc_port: EnvVar,
73    tracing_level: Option<EnvVar>,
74    typesense_api_key: EnvVar,
75    typesense_url: EnvVar,
76    web_port: EnvVar,
77}
78
79impl FloraConfig {
80    /// Initializes a new flora image config with the given parameters.
81    ///
82    /// Automatically generates an auth token.
83    #[must_use]
84    pub fn new(
85        auth_token: SecretString,
86        typesense_url: &Url,
87        typesense_api_key: SecretString,
88    ) -> Self {
89        let auth_token = EnvVar::Secret {
90            key: env_keys::AUTH_TOKEN.to_string(),
91            value: auth_token,
92        };
93
94        let typesense_api_key = EnvVar::Secret {
95            key: env_keys::TYPESENSE_API_KEY.to_string(),
96            value: typesense_api_key,
97        };
98        let typesense_url = EnvVar::Plain {
99            key: env_keys::TYPESENSE_URL.to_string(),
100            value: typesense_url.to_string(),
101        };
102
103        // Hardcode the ready message until we have a reason to let the caller configure it.
104        let ready_message = EnvVar::Plain {
105            key: env_keys::READY_MESSAGE_KEY.to_string(),
106            value: READY_MESSAGE.to_string(),
107        };
108
109        let db_path = Path::new(DATA_PATH_CONTAINER)
110            .join("flora.db")
111            .display()
112            .to_string();
113
114        // ditto with database URL.
115        let database_url = EnvVar::Plain {
116            key: env_keys::DATABASE_URL.into(),
117            value: db_path,
118        };
119
120        let rpc_port = EnvVar::Plain {
121            key: env_keys::SERVER_PORT_RPC.into(),
122            value: CONTAINER_PORT_RPC.to_string(),
123        };
124
125        let web_port = EnvVar::Plain {
126            key: env_keys::SERVER_PORT_WEB.into(),
127            value: CONTAINER_PORT_WEB.to_string(),
128        };
129
130        Self {
131            auth_token,
132            database_url,
133            ready_message,
134            rpc_port,
135            typesense_api_key,
136            typesense_url,
137            tracing_level: None,
138            web_port,
139        }
140    }
141
142    /// Sets the tracing level for the flora RPC server.
143    #[must_use]
144    pub fn with_tracing_level(mut self, tracing_level: Level) -> Self {
145        let tracing_level = EnvVar::Plain {
146            key: env_keys::TRACING_LEVEL.into(),
147            value: tracing_level.to_string(),
148        };
149        self.tracing_level = Some(tracing_level);
150        self
151    }
152
153    /// Returns the values in this config as a vec of environment variables.
154    fn into_vec(self) -> Vec<EnvVar> {
155        let mut vars = vec![
156            self.auth_token,
157            self.database_url,
158            self.ready_message,
159            self.rpc_port,
160            self.typesense_url,
161            self.typesense_api_key,
162            self.web_port,
163        ];
164
165        if let Some(tracing_level) = self.tracing_level {
166            vars.push(tracing_level);
167        }
168
169        vars
170    }
171}
172
173/// Omry server image for integration tests.
174#[allow(clippy::missing_docs_in_private_items)]
175#[derive(Debug)]
176pub struct FloraImage {
177    name: String,
178    tag: String,
179    env: Vec<EnvVar>,
180    mounts: Vec<Mount>,
181    container_ports: Vec<ContainerPort>,
182    ready_message: String,
183    host_ports: Option<HostPorts>,
184
185    // We need to store this on the struct so the temp directory
186    // sticks around until the struct is dropped.
187    #[allow(dead_code)]
188    data_temp_path: Option<TempPathWithString>,
189}
190
191impl FloraImage {
192    /// Creates a new flora image instance with the given flora server config and
193    /// Docker image config.
194    ///
195    /// ## Errors
196    /// Returns [`anyhow::Error`] if we can't create a temporary path for Omry's data.
197    pub fn new(flora_config: FloraConfig, image_config: ImageConfig) -> anyhow::Result<Self> {
198        let container_mount = ContainerMount::new(DATA_PATH_CONTAINER)?;
199        let container_ports = if image_config.should_expose_ports {
200            vec![
201                ContainerPort::Udp(CONTAINER_PORT_RPC),
202                ContainerPort::Tcp(CONTAINER_PORT_WEB),
203            ]
204        } else {
205            vec![]
206        };
207
208        let (_, ready_message) = flora_config.ready_message.as_key_value_pair();
209        let ready_message = ready_message.into();
210
211        Ok(Self {
212            name: IMAGE_NAME.to_string(),
213            tag: IMAGE_TAG.to_string(),
214            env: flora_config.into_vec(),
215            data_temp_path: container_mount.temp_path,
216            mounts: vec![container_mount.mount],
217            container_ports,
218            ready_message,
219            host_ports: None,
220        })
221    }
222
223    /// Configures this container to use the given ports on the host OS.
224    fn with_host_ports(mut self, ports: HostPorts) -> ContainerRequest<Self> {
225        self.host_ports = Some(ports);
226        self.with_mapped_port(ports.rpc, CONTAINER_PORT_RPC.udp())
227            .with_mapped_port(ports.web, CONTAINER_PORT_WEB.tcp())
228    }
229
230    /// Configures this container to request any free port from the host OS.
231    fn with_auto_host_ports(self) -> ContainerRequest<Self> {
232        let (host_ports, _dummy_listeners) = HostPorts::auto();
233        self.with_host_ports(host_ports)
234    }
235
236    /// Returns the host-side port this container is bound to (if any).
237    #[must_use]
238    pub fn host_ports(&self) -> Option<HostPorts> {
239        self.host_ports
240    }
241}
242
243impl Image for FloraImage {
244    fn name(&self) -> &str {
245        &self.name
246    }
247
248    fn tag(&self) -> &str {
249        &self.tag
250    }
251
252    fn ready_conditions(&self) -> Vec<WaitFor> {
253        vec![WaitFor::message_on_stdout(&self.ready_message)]
254    }
255
256    fn env_vars(
257        &self,
258    ) -> impl IntoIterator<Item = (impl Into<Cow<'_, str>>, impl Into<Cow<'_, str>>)> {
259        self.env.iter().map(EnvVar::as_key_value_pair)
260    }
261
262    fn mounts(&self) -> impl IntoIterator<Item = &Mount> {
263        &self.mounts
264    }
265
266    /// Returns the exposed ports for tarpc and web, in that order.
267    fn expose_ports(&self) -> &[ContainerPort] {
268        &self.container_ports
269    }
270}
271
272/// Type alias to reduce the type signature.
273pub type FloraContainer = ContainerAsync<FloraImage>;
274
275/// Starts an Omry RPC server container for integration tests.
276///
277/// ## Errors
278/// Returns [`anyhow::Error`] if we can't start the Omry server.
279pub async fn start_flora_container<L>(
280    flora_config: FloraConfig,
281    container_config: ContainerConfig<L>,
282) -> anyhow::Result<FloraContainer>
283where
284    L: ContainerLogConsumer,
285{
286    let flora_container =
287        FloraImage::new(flora_config, container_config.image_config)?.with_auto_host_ports();
288
289    let flora_container = container_config.configure_container(flora_container);
290    Ok(flora_container.start().await?)
291}