omnigraph/crates/omnigraph-cli/tests/support/mod.rs

#![allow(dead_code)]

use std::fs;
use std::net::TcpListener;
use std::path::{Path, PathBuf};
use std::process::{Child, Command as StdCommand, Output, Stdio};
use std::thread::sleep;
use std::time::Duration;

use assert_cmd::Command;
use reqwest::blocking::Client;
use serde_json::Value;
use tempfile::{TempDir, tempdir};

pub fn cli() -> Command {
    Command::cargo_bin("omnigraph").unwrap()
}

pub fn cli_process() -> StdCommand {
    StdCommand::new(assert_cmd::cargo::cargo_bin("omnigraph"))
}

fn server_process() -> StdCommand {
    if let Some(path) = std::env::var_os("CARGO_BIN_EXE_omnigraph-server") {
        StdCommand::new(path)
    } else if let Some(path) = built_server_binary() {
        StdCommand::new(path)
    } else {
        let cargo = std::env::var_os("CARGO").unwrap_or_else(|| "cargo".into());
        let mut cmd = StdCommand::new(cargo);
        cmd.arg("run")
            .arg("--quiet")
            .arg("-p")
            .arg("omnigraph-server")
            .arg("--");
        cmd
    }
}

fn built_server_binary() -> Option<PathBuf> {
    let workspace_root = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("../..");
    let candidate = workspace_root
        .join("target")
        .join("debug")
        .join(format!("omnigraph-server{}", std::env::consts::EXE_SUFFIX));
    candidate.exists().then_some(candidate)
}

pub fn fixture(name: &str) -> PathBuf {
    PathBuf::from(env!("CARGO_MANIFEST_DIR"))
        .join("../omnigraph/tests/fixtures")
        .join(name)
}

pub fn repo_path(root: &Path) -> PathBuf {
    root.join("demo.omni")
}

pub fn output_success(cmd: &mut Command) -> Output {
    let output = cmd.output().unwrap();
    assert!(
        output.status.success(),
        "command failed\nstdout:\n{}\nstderr:\n{}",
        String::from_utf8_lossy(&output.stdout),
        String::from_utf8_lossy(&output.stderr)
    );
    output
}

pub fn output_failure(cmd: &mut Command) -> Output {
    let output = cmd.output().unwrap();
    assert!(
        !output.status.success(),
        "command unexpectedly succeeded\nstdout:\n{}\nstderr:\n{}",
        String::from_utf8_lossy(&output.stdout),
        String::from_utf8_lossy(&output.stderr)
    );
    output
}

pub fn stdout_string(output: &Output) -> String {
    String::from_utf8(output.stdout.clone()).unwrap()
}

pub fn parse_stdout_json(output: &Output) -> Value {
    serde_json::from_slice(&output.stdout).unwrap()
}

pub fn init_repo(repo: &Path) {
    let schema = fixture("test.pg");
    output_success(cli().arg("init").arg("--schema").arg(&schema).arg(repo));
}

pub fn load_fixture(repo: &Path) {
    let data = fixture("test.jsonl");
    output_success(cli().arg("load").arg("--data").arg(&data).arg(repo));
}

pub fn write_jsonl(path: &Path, rows: &str) {
    fs::write(path, rows).unwrap();
}

pub fn write_query_file(path: &Path, source: &str) {
    fs::write(path, source).unwrap();
}

pub fn write_config(path: &Path, source: &str) {
    fs::write(path, source).unwrap();
}

pub fn write_file(path: &Path, source: &str) {
    fs::write(path, source).unwrap();
}

fn yaml_string(value: &str) -> String {
    format!("'{}'", value.replace('\'', "''"))
}

pub fn local_yaml_config(repo: &Path) -> String {
    format!(
        "\
graphs:
  local:
    uri: {}
cli:
  graph: local
  branch: main
query:
  roots:
    - .
policy: {{}}
",
        yaml_string(&repo.to_string_lossy())
    )
}

pub fn remote_yaml_config(url: &str) -> String {
    format!(
        "\
graphs:
  dev:
    uri: {}
cli:
  graph: dev
  branch: main
query:
  roots:
    - .
policy: {{}}
",
        yaml_string(url)
    )
}

pub struct TestServer {
    child: Child,
    pub base_url: String,
}

impl Drop for TestServer {
    fn drop(&mut self) {
        let _ = self.child.kill();
        let _ = self.child.wait();
    }
}

fn free_port() -> u16 {
    let listener = TcpListener::bind("127.0.0.1:0").unwrap();
    let port = listener.local_addr().unwrap().port();
    drop(listener);
    port
}

fn spawn_server_process(mut command: StdCommand) -> TestServer {
    let port = free_port();
    let bind = format!("127.0.0.1:{}", port);
    let mut child = command
        .arg("--bind")
        .arg(&bind)
        .stdout(Stdio::null())
        .stderr(Stdio::null())
        .spawn()
        .unwrap();
    let base_url = format!("http://{}", bind);
    let client = Client::new();
    for _ in 0..300 {
        if client
            .get(format!("{}/healthz", base_url))
            .send()
            .map(|response| response.status().is_success())
            .unwrap_or(false)
        {
            return TestServer { child, base_url };
        }
        if let Some(status) = child.try_wait().unwrap() {
            panic!("server exited before becoming healthy: {status}");
        }
        sleep(Duration::from_millis(100));
    }
    panic!("server did not become healthy");
}

pub fn spawn_server(repo: &Path) -> TestServer {
    let mut command = server_process();
    command.arg(repo);
    spawn_server_process(command)
}

pub fn spawn_server_with_config(config: &Path) -> TestServer {
    let mut command = server_process();
    command.arg("--config").arg(config);
    spawn_server_process(command)
}

pub fn spawn_server_with_config_env(config: &Path, envs: &[(&str, &str)]) -> TestServer {
    let mut command = server_process();
    command.arg("--config").arg(config);
    for (name, value) in envs {
        command.env(name, value);
    }
    spawn_server_process(command)
}


pub struct SystemRepo {
    _temp: TempDir,
    repo: PathBuf,
}

impl SystemRepo {
    pub fn initialized() -> Self {
        let temp = tempdir().unwrap();
        let repo = repo_path(temp.path());
        init_repo(&repo);
        Self { _temp: temp, repo }
    }

    pub fn loaded() -> Self {
        let temp = tempdir().unwrap();
        let repo = repo_path(temp.path());
        init_repo(&repo);
        load_fixture(&repo);
        Self { _temp: temp, repo }
    }

    pub fn path(&self) -> &Path {
        &self.repo
    }

    pub fn write_query(&self, name: &str, source: &str) -> PathBuf {
        let path = self.repo.parent().unwrap().join(name);
        write_query_file(&path, source);
        path
    }

    pub fn write_jsonl(&self, name: &str, rows: &str) -> PathBuf {
        let path = self.repo.parent().unwrap().join(name);
        write_jsonl(&path, rows);
        path
    }

    pub fn write_config(&self, name: &str, source: &str) -> PathBuf {
        let path = self.repo.parent().unwrap().join(name);
        write_config(&path, source);
        path
    }

    pub fn write_file(&self, name: &str, source: &str) -> PathBuf {
        let path = self.repo.parent().unwrap().join(name);
        write_file(&path, source);
        path
    }

    pub fn spawn_server(&self) -> TestServer {
        spawn_server(&self.repo)
    }

    pub fn spawn_server_with_config(&self, config: &Path) -> TestServer {
        spawn_server_with_config(config)
    }

    pub fn spawn_server_with_config_env(&self, config: &Path, envs: &[(&str, &str)]) -> TestServer {
        spawn_server_with_config_env(config, envs)
    }
}