use anyhow::{Context, Result};
use fluent_uri::Uri;
use omry_archiving::{Document, Marshal, Record, RecordParams, ToRecord};
use omry_service::iroh::EndpointAddr;
use std::env;
use std::fs::{self, File};
use std::io::prelude::*;
use std::net::{Ipv4Addr, Ipv6Addr, SocketAddrV4, SocketAddrV6, TcpListener, ToSocketAddrs};
use std::path::{Path, PathBuf};
use tempfile::{TempDir, tempdir};
use testcontainers::core::Mount;
use tracing::Level;
use tracing_subscriber::FmtSubscriber;
use url::Url;
use uuid::Uuid;
use crate::db::TempDb;
const DOCKER_HOST_KEY: &str = "DOCKER_HOST";
const WORKSPACE_DIR_KEY: &str = "CARGO_WORKSPACE_DIR";
const TEST_ARCHIVE_SUBPATH: &str = "test_data/example.postcard";
const FIND_PORT_MAX_RETRIES: u8 = 5;
fn free_port_from_socket_addr(
socket: impl ToSocketAddrs,
max_attempts: u8,
) -> anyhow::Result<(TcpListener, u16)> {
let mut attempts_left = max_attempts;
let socket_addrs = socket.to_socket_addrs()?.collect::<Vec<_>>();
while attempts_left > 0 {
if let Ok(listener) = TcpListener::bind(socket_addrs.as_slice()) {
let port = listener.local_addr()?.port();
return Ok((listener, port));
}
attempts_left -= 1;
}
anyhow::bail!("Couldn't bind a listener to a port in {max_attempts} attempts.");
}
#[must_use]
pub fn free_port_ipv4() -> (TcpListener, u16) {
let socket = SocketAddrV4::new(Ipv4Addr::LOCALHOST, 0);
#[allow(clippy::expect_used)]
free_port_from_socket_addr(socket, FIND_PORT_MAX_RETRIES).expect("Couldn't find a free port.")
}
#[must_use]
pub fn free_port_ipv6() -> (TcpListener, u16) {
let socket = SocketAddrV6::new(Ipv6Addr::LOCALHOST, 0, 0, 0);
#[allow(clippy::expect_used)]
free_port_from_socket_addr(socket, FIND_PORT_MAX_RETRIES).expect("Couldn't find a free port.")
}
#[derive(Debug, Clone)]
pub(crate) struct DockerHost {
uri: Uri<String>,
}
impl DockerHost {
pub(crate) fn from_env() -> Result<Self> {
let uri = env::var(DOCKER_HOST_KEY)
.map_err(|e| {
anyhow::anyhow!("Failed to get the docker host from the environment: {e}")
})?
.parse()
.map_err(|e| anyhow::anyhow!("Can't parse the docker host: {e}"))?;
Ok(Self { uri })
}
pub(crate) fn is_remote(&self) -> bool {
self.uri.scheme().as_str() == "tcp" && self.host().is_some()
}
pub(crate) fn host(&self) -> Option<&str> {
Some(self.uri.authority()?.host())
}
}
pub fn path_in_project(subpath: &str) -> Result<PathBuf> {
let (_, workspace_dir_path) = env::vars()
.find(|(k, _)| k == WORKSPACE_DIR_KEY)
.ok_or_else(|| anyhow::anyhow!("{WORKSPACE_DIR_KEY} is not set"))?;
let workspace_dir_path = fs::canonicalize(workspace_dir_path)?;
Ok(workspace_dir_path.join(subpath))
}
pub fn make_test_archive_holder() -> anyhow::Result<Document> {
let test_archive_path = path_in_project(TEST_ARCHIVE_SUBPATH)?;
let mut file = File::open(test_archive_path)?;
let mut buf = vec![];
file.read_to_end(&mut buf)?;
let archive_holder = Document::from_bytes(&buf)?;
Ok(archive_holder)
}
pub fn make_test_record_params() -> anyhow::Result<RecordParams> {
let document = make_test_archive_holder()?;
let url = "https://example.org/".to_string();
let title = "Example".to_string();
let client_datetime = "2024-10-11T13:49:46-05:00".to_string();
let language = Some("en".to_string());
Ok(RecordParams {
id: None,
url,
title,
language,
client_datetime,
document,
timestamp_flora: None,
updated_at: None,
})
}
pub fn make_test_record() -> anyhow::Result<Record> {
let incoming_record = Record::new(make_test_record_params()?);
Ok(incoming_record)
}
pub async fn make_test_db_records(how_many: usize) -> anyhow::Result<Vec<impl ToRecord>> {
let db = TempDb::new().await?;
let mut db_records = Vec::with_capacity(how_many);
for _ in 0..how_many {
let record = make_test_record()?;
let row_id = db.insert(record.try_into()?).await?;
let db_record = db.get_record(row_id).await?;
db_records.push(db_record);
}
Ok(db_records)
}
#[derive(Debug)]
pub struct TempPathWithString {
path: PathBuf,
path_string: String,
}
impl TempPathWithString {
pub fn new() -> anyhow::Result<Self> {
let temp_path = tempdir()?.keep();
#[allow(clippy::unnecessary_debug_formatting)]
let temp_path_string = temp_path
.to_str()
.ok_or_else(|| anyhow::anyhow!("Couldn't get a path string from {temp_path:?}"))?
.to_string();
Ok(Self {
path: temp_path,
path_string: temp_path_string,
})
}
#[must_use]
pub fn as_path(&self) -> &Path {
&self.path
}
#[must_use]
pub fn as_str(&self) -> &str {
&self.path_string
}
}
pub fn init_tracing_subscriber() -> anyhow::Result<()> {
let subscriber = FmtSubscriber::builder()
.with_max_level(Level::INFO)
.finish();
tracing::subscriber::set_global_default(subscriber)?;
Ok(())
}
#[derive(Debug)]
pub struct TempTracingLocation {
dir: TempDir,
dir_path: String,
file_prefix: String,
}
impl TempTracingLocation {
pub fn new() -> Result<Self> {
let dir = TempDir::new()?;
let dir_path = dir
.path()
.to_str()
.ok_or_else(|| anyhow::anyhow!("Cannot get the string path to the tracing dir."))?
.to_string();
let file_prefix = Uuid::now_v7().to_string();
Ok(Self {
dir,
dir_path,
file_prefix,
})
}
#[must_use]
pub fn dir(&self) -> &TempDir {
&self.dir
}
#[must_use]
pub fn path(&self) -> &str {
&self.dir_path
}
#[must_use]
pub fn file_prefix(&self) -> &str {
&self.file_prefix
}
pub fn read_first_trace(&self) -> Result<Option<String>> {
let dir = fs::read_dir(self.dir())?;
for entry in dir {
let entry = entry?;
let file_name = entry.file_name();
let file_name = file_name.to_string_lossy();
if file_name.starts_with(self.file_prefix()) {
let output = fs::read_to_string(entry.path())?;
return Ok(Some(output));
}
}
Ok(None)
}
}
#[derive(Debug, Copy, Clone)]
pub struct HostPorts {
pub rpc: u16,
pub web: u16,
}
impl HostPorts {
#[must_use]
pub(crate) fn auto() -> (Self, (TcpListener, TcpListener)) {
let (rpc_listener, rpc) = free_port_ipv6();
let (web_listener, web) = free_port_ipv6();
(Self { rpc, web }, (rpc_listener, web_listener))
}
}
#[derive(Debug)]
pub struct ContainerMount {
pub mount: Mount,
pub temp_path: Option<TempPathWithString>,
}
impl ContainerMount {
pub fn new(container_path: &str) -> Result<Self> {
if let Ok(docker_host) = DockerHost::from_env()
&& docker_host.is_remote()
{
let mount = Mount::tmpfs_mount(container_path);
return Ok(Self {
mount,
temp_path: None,
});
}
let temp_path = TempPathWithString::new()?;
let mount = Mount::bind_mount(temp_path.as_str(), container_path);
Ok(Self {
mount,
temp_path: Some(temp_path),
})
}
}
#[allow(clippy::missing_errors_doc)]
pub async fn endpoint_addr(server_root: &str) -> Result<EndpointAddr> {
const ENDPOINT_SUBPATH: &str = "/endpoint-addr";
let url = Url::parse(server_root).context("error parsing the URL")?;
let url = url.join(ENDPOINT_SUBPATH)?;
let bytes = reqwest::get(url)
.await?
.error_for_status()
.context("Couldn't get the endpoint_address")?
.bytes()
.await
.context("Couldn't read the bytes")?;
let endpoint_addr =
postcard::from_bytes(&bytes).context("Error deserializing the endpoint address")?;
Ok(endpoint_addr)
}
pub mod token {
use rand::RngExt;
use rand::distr::Alphanumeric;
use rand::rngs::ChaCha20Rng;
pub fn generate(len: usize) -> String {
let rng = rand::make_rng::<ChaCha20Rng>();
rng.sample_iter(Alphanumeric)
.take(len)
.map(char::from)
.collect()
}
}