omry_integration/
helpers.rs

1//! Helper functions and data structures for integration tests.
2use anyhow::{Context, Result};
3use fluent_uri::Uri;
4use omry_archiving::{Document, Marshal, Record, RecordParams, ToRecord};
5use omry_service::iroh::EndpointAddr;
6use std::env;
7use std::fs::{self, File};
8use std::io::prelude::*;
9use std::net::{Ipv4Addr, Ipv6Addr, SocketAddrV4, SocketAddrV6, TcpListener, ToSocketAddrs};
10use std::path::{Path, PathBuf};
11use tempfile::{TempDir, tempdir};
12use testcontainers::core::Mount;
13use tracing::Level;
14use tracing_subscriber::FmtSubscriber;
15use url::Url;
16use uuid::Uuid;
17
18use crate::db::TempDb;
19
20// Env var key for the Docker host.
21const DOCKER_HOST_KEY: &str = "DOCKER_HOST";
22
23/// Environment variable key that corresponds to the Cargo workspace directory.
24const WORKSPACE_DIR_KEY: &str = "CARGO_WORKSPACE_DIR";
25
26/// The path to the serialized web page used in tests.
27const TEST_ARCHIVE_SUBPATH: &str = "test_data/example.postcard";
28
29/// The path to an HTML file for a long article,
30/// that is expected to be split into many search chunks.
31const TEST_LONG_DOC_SUBPATH: &str = "test_data/lipsum.html";
32
33/// Maximum number of times we'll try to find a free port to bind to.
34const FIND_PORT_MAX_RETRIES: u8 = 5;
35
36/// Gets a free port for the given socket from the OS. Returns the listener and the port.
37///
38/// The listener is returned to deal with TOCTOU, so that the OS
39/// doesn't give the port to another service before we bind to it for real
40/// wherever we actually need that port in our tests.
41fn free_port_from_socket_addr(
42    socket: impl ToSocketAddrs,
43    max_attempts: u8,
44) -> anyhow::Result<(TcpListener, u16)> {
45    let mut attempts_left = max_attempts;
46    let socket_addrs = socket.to_socket_addrs()?.collect::<Vec<_>>();
47
48    while attempts_left > 0 {
49        if let Ok(listener) = TcpListener::bind(socket_addrs.as_slice()) {
50            let port = listener.local_addr()?.port();
51            return Ok((listener, port));
52        }
53        attempts_left -= 1;
54    }
55
56    anyhow::bail!("Couldn't bind a listener to a port in {max_attempts} attempts.");
57}
58
59/// Gets a free IPv4 port from the OS. Returns the listener and the port.
60///
61/// The listener is returned to account for TOCTOU, so that the OS
62/// doesn't release the port for another service before we bind to it for real
63/// in our tests.
64///
65/// ## Panics
66/// If the port can't be found within some number of retries (currently 3).
67#[must_use]
68pub fn free_port_ipv4() -> (TcpListener, u16) {
69    let socket = SocketAddrV4::new(Ipv4Addr::LOCALHOST, 0);
70
71    // Okay to panic if we can't find a free port, since we can't run integration tests
72    // without a port anyway.
73    #[allow(clippy::expect_used)]
74    free_port_from_socket_addr(socket, FIND_PORT_MAX_RETRIES).expect("Couldn't find a free port.")
75}
76
77/// Gets a free IPv6 port from the OS. See [`free_port_ipv4`] for details about the return value.
78///
79/// ## Panics
80/// If the port can't be found within some number of retries (currently 3).
81#[must_use]
82pub fn free_port_ipv6() -> (TcpListener, u16) {
83    let socket = SocketAddrV6::new(Ipv6Addr::LOCALHOST, 0, 0, 0);
84
85    // Okay to panic if we can't find a free port, since we can't run integration tests
86    // without a port anyway.
87    #[allow(clippy::expect_used)]
88    free_port_from_socket_addr(socket, FIND_PORT_MAX_RETRIES).expect("Couldn't find a free port.")
89}
90
91/// The Docker host (where the Docker socket is running).
92#[derive(Debug, Clone)]
93pub(crate) struct DockerHost {
94    uri: Uri<String>,
95}
96
97impl DockerHost {
98    /// Gets the URI for the `DOCKER_HOST`, if any.
99    ///
100    ///
101    /// # Errors
102    ///
103    /// Returns an error if:
104    /// - `DOCKER_HOST` is not set
105    /// - `DOCKER_HOST` cannot be parsed into a [`Uri`].
106    pub(crate) fn from_env() -> Result<Self> {
107        let uri = env::var(DOCKER_HOST_KEY)
108            .map_err(|e| {
109                anyhow::anyhow!("Failed to get the docker host from the environment: {e}")
110            })?
111            .parse()
112            .map_err(|e| anyhow::anyhow!("Can't parse the docker host: {e}"))?;
113
114        Ok(Self { uri })
115    }
116
117    /// Returns [`true`] if the `DOCKER_HOST` appears to be remote.
118    ///
119    /// The heuristic is (all must be true):
120    ///
121    /// - `DOCKER_HOST` env var is set,
122    /// - `DOCKER_HOST` URI scheme is `tcp`,
123    /// - `DOCKER_HOST` host is pretent
124    pub(crate) fn is_remote(&self) -> bool {
125        self.uri.scheme().as_str() == "tcp" && self.host().is_some()
126    }
127
128    /// Gets the host component of the Docker host (if any).
129    ///
130    /// Will only be present if the Docker host is remote.
131    pub(crate) fn host(&self) -> Option<&str> {
132        Some(self.uri.authority()?.host())
133    }
134}
135
136/// Takes the given subpath, and returns the canonical path within
137/// the current Cargo project.
138///
139/// ## Errors
140/// Returns [`anyhow::Error`] if we fail to:
141///
142/// - Get the path to the current project via `CARGO_WORKSPACE_DIR` env variable
143/// - Get the full canonical path from it
144pub fn path_in_project(subpath: &str) -> Result<PathBuf> {
145    let (_, workspace_dir_path) = env::vars()
146        .find(|(k, _)| k == WORKSPACE_DIR_KEY)
147        .ok_or_else(|| anyhow::anyhow!("{WORKSPACE_DIR_KEY} is not set"))?;
148
149    let workspace_dir_path = fs::canonicalize(workspace_dir_path)?;
150    Ok(workspace_dir_path.join(subpath))
151}
152
153/// Loads a document to be used in tests.
154///
155/// ## Errors
156/// Returns [`anyhow::Error`] if we fail to:
157///
158/// - Get the path to the current project via `CARGO_WORKSPACE_DIR` env variable
159/// - Get the full canonical path from it
160/// - Read the file from that path
161/// - Deserialize the data from that file
162pub fn make_test_archive_holder() -> anyhow::Result<Document> {
163    let test_archive_path = path_in_project(TEST_ARCHIVE_SUBPATH)?;
164    let mut file = File::open(test_archive_path)?;
165    let mut buf = vec![];
166    file.read_to_end(&mut buf)?;
167    let archive_holder = Document::from_bytes(&buf)?;
168    Ok(archive_holder)
169}
170
171/// Makes an instance of [`RecordParams`] for tests.
172///
173/// ## Errors
174/// Returns [`anyhow::Error`] if the test document can't be loaded.
175pub fn make_test_record_params() -> anyhow::Result<RecordParams> {
176    let document = make_test_archive_holder()?;
177    let url = "https://example.org/".to_string();
178    let title = "Example".to_string();
179    let client_datetime = "2024-10-11T13:49:46-05:00".to_string();
180    let language = Some("en".to_string());
181
182    Ok(RecordParams {
183        id: None,
184        url,
185        title,
186        language,
187        client_datetime,
188        document,
189        timestamp_flora: None,
190        updated_at: None,
191    })
192}
193
194/// Makes a [`Record`] instance for tests.
195///
196/// ## Errors
197/// Returns [`anyhow::Error`] if the test record can't be loaded.
198pub fn make_test_record() -> anyhow::Result<Record> {
199    let incoming_record = Record::new(make_test_record_params()?);
200    Ok(incoming_record)
201}
202
203fn load_long_doc() -> anyhow::Result<Document> {
204    let (_, workspace_dir_path) = env::vars()
205        .find(|(k, _)| k == WORKSPACE_DIR_KEY)
206        .ok_or_else(|| anyhow::anyhow!("{WORKSPACE_DIR_KEY} is not set"))?;
207    let workspace_dir_path = fs::canonicalize(workspace_dir_path)?;
208    let test_long_doc_path = workspace_dir_path.join(TEST_LONG_DOC_SUBPATH);
209
210    let mut file = File::open(&test_long_doc_path)?;
211    let mut buf = vec![];
212    file.read_to_end(&mut buf)?;
213
214    let archive_holder = Document::new_web_page_unchecked(buf);
215    Ok(archive_holder)
216}
217
218/// Creates record params with a 'long' document.
219///
220/// # Errors
221/// Returns [`anyhow::Error`] if the test record can't be loaded.
222pub fn make_long_doc_record_params() -> anyhow::Result<RecordParams> {
223    let document = load_long_doc()?;
224    let url = "https://www.lipsum.com".to_string();
225    let title = "Lorem Ipsum";
226    let client_datetime = "2026-09-02T12:00:00-05:00";
227
228    #[allow(clippy::unreadable_literal)]
229    let timestamp_flora = Some(1788371331);
230    let language = Some("en".to_string());
231
232    Ok(RecordParams {
233        id: None,
234        url,
235        title: title.to_string(),
236        client_datetime: client_datetime.to_string(),
237        timestamp_flora,
238        language,
239        document,
240        updated_at: None,
241    })
242}
243
244/// Gets records that have been inserted into the in-memory SQLite,
245/// so that they have unique ids.
246///
247/// ## Errors
248/// Returns [`anyhow::Error`] if any of the following operations fail:
249///
250/// - We can't create a new [`TempDb`].
251/// - We can't insert the test record into SQLite.
252/// - We can't get that record back from SQLite.
253pub async fn make_test_db_records(how_many: usize) -> anyhow::Result<Vec<impl ToRecord>> {
254    let db = TempDb::new().await?;
255    let mut db_records = Vec::with_capacity(how_many);
256    for _ in 0..how_many {
257        let record = make_test_record()?;
258        let row_id = db.insert(record.try_into()?).await?;
259        let db_record = db.get_record(row_id).await?;
260        db_records.push(db_record);
261    }
262    Ok(db_records)
263}
264
265/// Temporary path and its string representation.
266#[derive(Debug)]
267pub struct TempPathWithString {
268    /// The path.
269    path: PathBuf,
270
271    /// The string representation of the same path, to make it easier to use
272    /// in some functions.
273    path_string: String,
274}
275
276impl TempPathWithString {
277    /// Creates a new [`tempfile::TempDir`], and stores its associated path
278    /// and string representation.
279    ///
280    /// ## Errors
281    /// If we can't create a temporary directory or construct the string representation
282    /// of its path.
283    pub fn new() -> anyhow::Result<Self> {
284        let temp_path = tempdir()?.keep();
285
286        // If we can't get a string representation of the temp path, then its display
287        // repesentation is unlikely to be useful.
288        // So clippy's suggestion isn't helpful here.
289        #[allow(clippy::unnecessary_debug_formatting)]
290        let temp_path_string = temp_path
291            .to_str()
292            .ok_or_else(|| anyhow::anyhow!("Couldn't get a path string from {temp_path:?}"))?
293            .to_string();
294
295        Ok(Self {
296            path: temp_path,
297            path_string: temp_path_string,
298        })
299    }
300
301    /// Returns a reference to this temporary directory's path.
302    #[must_use]
303    pub fn as_path(&self) -> &Path {
304        &self.path
305    }
306
307    /// Returns a reference to this temporary path's string representation.
308    #[must_use]
309    pub fn as_str(&self) -> &str {
310        &self.path_string
311    }
312}
313
314/// Initializes a tracing subscriber for tests.
315///
316/// ## Errors
317/// Returns [`anyhow::Error`] if
318/// we can't set the tracing subscriber as the global default.
319pub fn init_tracing_subscriber() -> anyhow::Result<()> {
320    let subscriber = FmtSubscriber::builder()
321        .with_max_level(Level::INFO)
322        .finish();
323    tracing::subscriber::set_global_default(subscriber)?;
324
325    Ok(())
326}
327
328/// A temporary location for storing traces.
329///
330/// Used for tests that need to verify certain events happened
331/// based on the application traces.
332#[derive(Debug)]
333pub struct TempTracingLocation {
334    /// The temporary directory to store the trace.
335    dir: TempDir,
336
337    /// The string representation of the temp tracing path, pre-computed
338    /// in [`TempTracingLocation::new`].
339    /// Storing the pre-computed string makes it easier for users of this
340    /// structure to get a string path without having to do their own error handling.
341    dir_path: String,
342
343    /// The UUID7 prefix of the trace files in the temporary directory.
344    file_prefix: String,
345}
346
347impl TempTracingLocation {
348    /// Creates a new temporary directory for trace files, and generates their
349    /// UUID7 filename prefix.
350    ///
351    /// ## Errors
352    /// Returns [`anyhow::Error`] if we fail to:
353    /// - Create a new temporary directory
354    /// - Get the string representation of that directory's path
355    pub fn new() -> Result<Self> {
356        let dir = TempDir::new()?;
357        let dir_path = dir
358            .path()
359            .to_str()
360            .ok_or_else(|| anyhow::anyhow!("Cannot get the string path to the tracing dir."))?
361            .to_string();
362
363        let file_prefix = Uuid::now_v7().to_string();
364
365        Ok(Self {
366            dir,
367            dir_path,
368            file_prefix,
369        })
370    }
371
372    /// Returns a reference to the temporary directory which will hold the trace files.
373    #[must_use]
374    pub fn dir(&self) -> &TempDir {
375        &self.dir
376    }
377
378    /// Returns a string representation of the path, corresponding to [`Self::dir`].
379    #[must_use]
380    pub fn path(&self) -> &str {
381        &self.dir_path
382    }
383
384    /// The trace filename prefix.
385    #[must_use]
386    pub fn file_prefix(&self) -> &str {
387        &self.file_prefix
388    }
389
390    /// Iterates over the temporary directory and finds the first trace file that matches the [`Self::file_prefix`].
391    ///
392    /// ## Returns
393    /// The contents of the found trace file, if any.
394    ///
395    /// ## Errors
396    /// Returns [`anyhow::Error`] if we fail to:
397    /// - Read the temporary directory
398    /// - Get entries in that directory
399    /// - Read the entry matching the given prefix into a string
400    pub fn read_first_trace(&self) -> Result<Option<String>> {
401        let dir = fs::read_dir(self.dir())?;
402
403        for entry in dir {
404            let entry = entry?;
405            let file_name = entry.file_name();
406            let file_name = file_name.to_string_lossy();
407
408            if file_name.starts_with(self.file_prefix()) {
409                let output = fs::read_to_string(entry.path())?;
410                return Ok(Some(output));
411            }
412        }
413
414        Ok(None)
415    }
416}
417
418/// Host-side ports for the Omry server container to use.
419#[derive(Debug, Copy, Clone)]
420pub struct HostPorts {
421    /// Port for tarpc to bind to.
422    pub rpc: u16,
423
424    /// Port for the web server to bind to.
425    pub web: u16,
426}
427
428impl HostPorts {
429    /// Requests any free ports from the host OS.
430    ///
431    /// Returns dummy listeners for both respective ports to avoid
432    /// the ports becoming unavailable due to TOCTOU before
433    /// the caller can bind to them.
434    #[must_use]
435    pub(crate) fn auto() -> (Self, (TcpListener, TcpListener)) {
436        let (rpc_listener, rpc) = free_port_ipv6();
437        let (web_listener, web) = free_port_ipv6();
438        (Self { rpc, web }, (rpc_listener, web_listener))
439    }
440}
441
442/// A mount for a container.
443#[derive(Debug)]
444pub struct ContainerMount {
445    /// The mount for the container.
446    pub mount: Mount,
447
448    /// Temporary path on the host. Used if the Docker host is local.
449    pub temp_path: Option<TempPathWithString>,
450}
451
452impl ContainerMount {
453    /// Creates a mount for the container with the given path.
454    ///
455    /// If we're using a remote Docker host (set by `USING_REMOTE_DOCKER_HOST_KEY`),
456    /// will use [`Mount::tmpfs_mount`].
457    /// This assumes the remote Docker host is running on Linux (and will fail otherwise).
458    ///
459    /// Otherwise, will use [`Mount::bind_mount`].
460    ///
461    /// # Errors
462    ///
463    /// Returns [`anyhow::Error`] if the Docker host is local,
464    /// and we can't create a temporary path on the host.
465    pub fn new(container_path: &str) -> Result<Self> {
466        if let Ok(docker_host) = DockerHost::from_env()
467            && docker_host.is_remote()
468        {
469            let mount = Mount::tmpfs_mount(container_path);
470            return Ok(Self {
471                mount,
472                temp_path: None,
473            });
474        }
475
476        let temp_path = TempPathWithString::new()?;
477        let mount = Mount::bind_mount(temp_path.as_str(), container_path);
478
479        Ok(Self {
480            mount,
481            temp_path: Some(temp_path),
482        })
483    }
484}
485
486/// Gets the iroh endpoint address for the given server,
487/// via its web API.
488#[allow(clippy::missing_errors_doc)]
489pub async fn endpoint_addr(server_root: &str) -> Result<EndpointAddr> {
490    const ENDPOINT_SUBPATH: &str = "/endpoint-addr";
491    let url = Url::parse(server_root).context("error parsing the URL")?;
492    let url = url.join(ENDPOINT_SUBPATH)?;
493    let bytes = reqwest::get(url)
494        .await?
495        .error_for_status()
496        .context("Couldn't get the endpoint_address")?
497        .bytes()
498        .await
499        .context("Couldn't read the bytes")?;
500    let endpoint_addr =
501        postcard::from_bytes(&bytes).context("Error deserializing the endpoint address")?;
502    Ok(endpoint_addr)
503}
504
505/// Module for generating an authentication token.
506///
507/// Inspired by
508/// <https://github.com/EffortlessMetrics/uselesskey/blob/52375d7e2987de345006ce50e7831383ddbc0ab4/crates/uselesskey-core-token-shape/src/lib.rs#L70>.
509pub mod token {
510    use rand::RngExt;
511    use rand::distr::Alphanumeric;
512    use rand::rngs::ChaCha20Rng;
513
514    /// Generates a new random token with the given length.
515    pub fn generate(len: usize) -> String {
516        let rng = rand::make_rng::<ChaCha20Rng>();
517        rng.sample_iter(Alphanumeric)
518            .take(len)
519            .map(char::from)
520            .collect()
521    }
522}
omry_integration/helpers.rs

omry_integration/
helpers.rs
