typedb_test/

lib.rs

#![doc = include_str!("../README.md")]
#[allow(unused_imports)]
#[macro_use]
pub extern crate googletest;

#[cfg(test)]
mod tests;

pub use googletest::gtest;
use {
    googletest::prelude::{Result, StaticFixture},
    port_check::free_local_port_in_range,
    std::{
        fmt::Display,
        ops::Deref,
        panic::RefUnwindSafe,
        path::{Path, PathBuf},
        process::{Child, Command},
        sync::{Arc, Mutex},
        time::Duration,
    },
    typedb_driver::{Credentials, DriverOptions, TypeDBDriver},
};

const DEFAULT_DATABASE_NAME: &str = "test";

#[derive(Debug, Clone)]
pub struct TypeDBFixture {
    driver: Arc<typedb_driver::TypeDBDriver>,
    child: Arc<Mutex<Child>>,
}

impl Drop for TypeDBFixture {
    fn drop(&mut self) {
        self.child.lock().unwrap().kill().unwrap();
    }
}

impl RefUnwindSafe for TypeDBFixture {}

#[derive(Debug, Clone, Default)]
enum OS {
    #[cfg_attr(any(target_os = "linux"), default)]
    #[cfg_attr(not(target_os = "linux"), allow(dead_code))]
    Linux,
    #[cfg_attr(any(target_os = "macos"), default)]
    #[cfg_attr(not(target_os = "macos"), allow(dead_code))]
    Mac,
    #[cfg_attr(any(target_os = "windows"), default)]
    #[cfg_attr(not(target_os = "windows"), allow(dead_code))]
    Windows,
    #[cfg_attr(
        not(any(target_os = "linux", target_os = "macos", target_os = "windows")),
        default
    )]
    #[allow(dead_code)]
    Unknown,
}

#[derive(Debug, Clone, Default)]
enum Architecture {
    #[cfg_attr(any(target_arch = "x86_64"), default)]
    #[cfg_attr(not(target_arch = "x86_64"), allow(dead_code))]
    X86_64,
    #[cfg_attr(any(target_arch = "aarch64"), default)]
    #[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
    Aarch64,
    #[cfg_attr(not(any(target_arch = "x86_64", target_arch = "aarch64")), default)]
    #[allow(dead_code)]
    Unknown,
}

impl Display for OS {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            OS::Linux => write!(f, "linux"),
            OS::Mac => write!(f, "mac"),
            OS::Windows => write!(f, "windows"),
            OS::Unknown => unimplemented!(),
        }
    }
}

impl Display for Architecture {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Architecture::X86_64 => write!(f, "x86_64"),
            Architecture::Aarch64 => write!(f, "arm64"),
            Architecture::Unknown => unimplemented!(),
        }
    }
}

#[allow(dead_code)]
fn default_platform() -> String {
    format!("{}-{}", OS::default(), Architecture::default())
}

impl TypeDBFixture {
    #[allow(dead_code)]
    fn download_and_extract_typedb(archive_path: &Path, extract_dir: &Path) -> Result<()> {
        // Download TypeDB using curl
        let output = Command::new("curl")
            .args([
                "-L",
                &format!(
                    "https://repo.typedb.com/public/public-release/raw/names/typedb-all-{}/versions/latest/download",
                    default_platform()
                ),
                "-o",
                archive_path.to_str().unwrap(),
            ])
            .output()?;

        if !output.status.success() {
            panic!("Failed to extract TypeDB: {}", unsafe {
                String::from_utf8_unchecked(output.stderr)
            });
        }

        // Extract TypeDB
        let output = Command::new("tar")
            .args([
                "-xzf",
                archive_path.to_str().unwrap(),
                "-C",
                extract_dir.to_str().unwrap(),
            ])
            .output()?;

        if !output.status.success() {
            panic!("Failed to extract TypeDB: {}", unsafe {
                String::from_utf8_unchecked(output.stderr)
            });
        }

        Ok(())
    }

    async fn start_typedb_server(typedb_dir: &Path, port: u16) -> Result<Child> {
        let typedb_binary = typedb_dir.join("typedb");

        let child = Command::new(&typedb_binary)
            .args(["server", "--server.address", &format!("localhost:{}", port)])
            .spawn()?;

        // Give the server time to start
        tokio::time::sleep(Duration::from_secs(5)).await;

        Ok(child)
    }
}

impl StaticFixture for TypeDBFixture {
    fn set_up_once() -> googletest::Result<Self> {
        let temp_dir = tempdir::TempDir::new("typedb_test")?;
        let typedb_dir = temp_dir.into_path().to_path_buf();
        let archive_path: PathBuf = typedb_dir.join("typedb.tar.gz");
        let extracted_dir: PathBuf = typedb_dir;

        Self::download_and_extract_typedb(&archive_path, &extracted_dir).unwrap();
        let free_port_in_range = free_local_port_in_range(10000..=15000).unwrap();

        let address = format!("localhost:{}", free_port_in_range);
        let credentials = Credentials::new("admin", "password");
        let options = DriverOptions::new(false, None).unwrap();

        let mut driver: Option<TypeDBDriver> = None;
        let mut child: Option<Child> = None;

        // binary path {extracted_dir}/typedb-all-{default_platform()}-{version}/typedb
        // the version can be determined after extraction.
        // it'll be the first directory in the extracted directory
        let typedb_binary_path = extracted_dir.read_dir().unwrap();
        let mut dir = None;
        for entry in typedb_binary_path.flatten() {
            if entry.file_type().unwrap().is_dir() {
                dir = Some(entry.path());
                break;
            }
        }
        if dir.is_none() {
            panic!("No directory found in extracted directory");
        }
        let dir = dir.unwrap();
        let rt = tokio::runtime::Runtime::new().unwrap();
        rt.block_on(async {
            let server_child = Self::start_typedb_server(&dir, free_port_in_range)
                .await
                .unwrap();
            child = Some(server_child);
            let _driver = TypeDBDriver::new(address, credentials, options)
                .await
                .unwrap();
            _driver
                .databases()
                .create(DEFAULT_DATABASE_NAME)
                .await
                .unwrap();
            driver = Some(_driver);
        });
        Ok(TypeDBFixture {
            driver: driver.take().unwrap().into(),
            child: Arc::new(Mutex::new(child.take().unwrap())),
        })
    }
}

impl Deref for TypeDBFixture {
    type Target = TypeDBDriver;

    fn deref(&self) -> &Self::Target {
        &self.driver
    }
}