test: add golden test for tabby server (#566)

* test: add golden test for tabby server

* fix

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"

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

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"
+                }
+            ]
+        }
+    }
+]

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<TestCase> = serdeconv::from_json_file(golden_path()).unwrap();
+    for case in cases {
+        golden_test(case.request, case.expected).await;
+    }
+}