From 1503ef6abaf4b4a97173e603aabfbc1486cfba4f Mon Sep 17 00:00:00 2001 From: Meng Zhang Date: Mon, 16 Oct 2023 01:06:37 -0700 Subject: [PATCH] test: add golden test for tabby server (#566) * test: add golden test for tabby server * fix --- Cargo.lock | 61 +++++++++++++++++++ crates/tabby/Cargo.toml | 4 ++ crates/tabby/tests/golden.json | 35 +++++++++++ crates/tabby/tests/goldentests.rs | 97 +++++++++++++++++++++++++++++++ 4 files changed, 197 insertions(+) create mode 100644 crates/tabby/tests/golden.json create mode 100644 crates/tabby/tests/goldentests.rs diff --git a/Cargo.lock b/Cargo.lock index 63d94f5..6cffadc 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -148,6 +148,16 @@ version = "1.6.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "bddcadddf5e9015d310179a59bb28c4d4b9920ad0f11e8e14dbadf654890c9a6" +[[package]] +name = "assert-json-diff" +version = "2.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "47e4f2b81832e72834d7518d8487a0396a28cc408186a2e8854c0f98011faf12" +dependencies = [ + "serde", + "serde_json", +] + [[package]] name = "async-stream" version = "0.3.5" @@ -859,6 +869,12 @@ dependencies = [ "syn 1.0.109", ] +[[package]] +name = "diff" +version = "0.1.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "56254986775e3233ffa9c4d7d3faaf6d36a2c09d30b20687e9f88bc8bafc16c8" + [[package]] name = "digest" version = "0.10.7" @@ -1203,6 +1219,21 @@ version = "0.3.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d2fabcfbdc87f4758337ca535fb41a6d701b65693ce38287d856d1674551ec9b" +[[package]] +name = "goldie" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fd1d4b95ae93c6d91591a2998aa7363113e51130ede293b3c92ac89b63e13914" +dependencies = [ + "anyhow", + "once_cell", + "pretty_assertions", + "serde", + "serde_json", + "upon", + "yansi", +] + [[package]] name = "h2" version = "0.3.19" @@ -2304,6 +2335,16 @@ version = "0.2.17" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5b40af805b3121feab8a3c29f04d8ad262fa8e0561883e7653e024ae4479e6de" +[[package]] +name = "pretty_assertions" +version = "1.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "af7cee1a6c8a5b9208b3cb1061f10c0cb689087b3d8ce85fb9d2dd7a29b6ba66" +dependencies = [ + "diff", + "yansi", +] + [[package]] name = "prettyplease" version = "0.1.25" @@ -3097,6 +3138,7 @@ name = "tabby" version = "0.4.0-dev" dependencies = [ "anyhow", + "assert-json-diff", "async-stream", "axum", "axum-streams", @@ -3104,6 +3146,7 @@ dependencies = [ "clap", "ctranslate2-bindings", "futures", + "goldie", "http-api-bindings", "hyper", "lazy_static", @@ -3114,6 +3157,7 @@ dependencies = [ "opentelemetry", "opentelemetry-otlp", "regex", + "reqwest", "rust-embed 8.0.0", "serde", "serde_json", @@ -4049,6 +4093,17 @@ version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39ec24b3121d976906ece63c9daad25b85969647682eee313cb5779fdd69e14e" +[[package]] +name = "upon" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b77ce40602cb1a7dfcdd6336f6d8baa2803c898aafbc0d46156b59727f2e7135" +dependencies = [ + "serde", + "unicode-ident", + "unicode-width", +] + [[package]] name = "url" version = "2.3.1" @@ -4537,6 +4592,12 @@ dependencies = [ "libc", ] +[[package]] +name = "yansi" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09041cd90cf85f7f8b2df60c646f853b7f535ce68f85244eb6731cf89fa498ec" + [[package]] name = "zip" version = "0.6.6" diff --git a/crates/tabby/Cargo.toml b/crates/tabby/Cargo.toml index 2076f73..24f4c02 100644 --- a/crates/tabby/Cargo.toml +++ b/crates/tabby/Cargo.toml @@ -60,3 +60,7 @@ link_shared = ["ctranslate2-bindings/link_shared"] [build-dependencies] vergen = { version = "8.0.0", features = ["build", "git", "gitcl"] } + +[dev-dependencies] +assert-json-diff = "2.0.2" +reqwest.workspace = true diff --git a/crates/tabby/tests/golden.json b/crates/tabby/tests/golden.json new file mode 100644 index 0000000..7f968d2 --- /dev/null +++ b/crates/tabby/tests/golden.json @@ -0,0 +1,35 @@ +[ + { + "request": { + "language": "python", + "segments": { + "prefix": "def fib(n):\n ", + "suffix": "\n return fib(n - 1) + fib(n - 2)" + } + }, + "expected": { + "choices": [ + { + "index": 0, + "text": " if n <= 1:\n return n" + } + ] + } + }, + { + "request": { + "language": "python", + "segments": { + "prefix": "import datetime\n\ndef parse_expenses(expenses_string):\n \"\"\"Parse the list of expenses and return the list of triples (date, value, currency).\n Ignore lines starting with #.\n Parse the date using datetime.\n Example expenses_string:\n 2016-01-02 -34.01 USD\n 2016-01-03 2.59 DKK\n 2016-01-03 -2.72 EUR\n \"\"\"\n for line in expenses_string.split('\\n'):\n " + } + }, + "expected": { + "choices": [ + { + "index": 0, + "text": "if line.startswith('#'):\n continue\n date, value, currency = line.split()\n date = datetime.datetime.strptime(date, '%Y-%m-%d')\n yield date, float(value), currency" + } + ] + } + } +] \ No newline at end of file diff --git a/crates/tabby/tests/goldentests.rs b/crates/tabby/tests/goldentests.rs new file mode 100644 index 0000000..098cfb6 --- /dev/null +++ b/crates/tabby/tests/goldentests.rs @@ -0,0 +1,97 @@ +use std::path::PathBuf; + +use assert_json_diff::assert_json_include; +use lazy_static::lazy_static; +use serde::Deserialize; +use tokio::{ + process::Command, + time::{sleep, Duration}, +}; + +lazy_static! { + static ref SERVER: bool = { + let mut cmd = Command::new(tabby_path()); + cmd.arg("serve") + .arg("--model") + .arg("TabbyML/StarCoder-1B") + .arg("--port") + .arg("9090") + .arg("--device") + .arg("metal") + .kill_on_drop(true); + tokio::task::spawn(async move { + cmd.spawn() + .expect("Failed to start server") + .wait() + .await + .unwrap(); + }); + true + }; + static ref CLIENT: reqwest::Client = reqwest::Client::new(); +} + +fn workspace_dir() -> PathBuf { + let output = std::process::Command::new(env!("CARGO")) + .arg("locate-project") + .arg("--workspace") + .arg("--message-format=plain") + .output() + .unwrap() + .stdout; + let cargo_path = std::path::Path::new(std::str::from_utf8(&output).unwrap().trim()); + cargo_path.parent().unwrap().to_path_buf() +} + +fn tabby_path() -> PathBuf { + workspace_dir().join("target/debug/tabby") +} + +fn golden_path() -> PathBuf { + workspace_dir().join("crates/tabby/tests/golden.json") +} + +async fn wait_for_server() { + lazy_static::initialize(&SERVER); + + loop { + println!("Waiting for server to start..."); + let is_ok = reqwest::get("http://localhost:9090/v1/health") + .await + .is_ok(); + if is_ok { + break; + } else { + sleep(Duration::from_secs(5)).await; + } + } +} + +async fn golden_test(body: serde_json::Value, expected: serde_json::Value) { + let actual: serde_json::Value = CLIENT + .post("http://localhost:9090/v1/completions") + .json(&body) + .send() + .await + .unwrap() + .json() + .await + .unwrap(); + assert_json_include!(actual: actual, expected: expected); +} + +#[derive(Deserialize)] +struct TestCase { + request: serde_json::Value, + expected: serde_json::Value, +} + +#[tokio::test] +async fn run_golden_tests() { + wait_for_server().await; + + let cases: Vec = serdeconv::from_json_file(golden_path()).unwrap(); + for case in cases { + golden_test(case.request, case.expected).await; + } +}