From 614d993344bb4e09909631204adf7fd9b3478632 Mon Sep 17 00:00:00 2001
From: Meng Zhang <meng@tabbyml.com>
Date: Tue, 4 Apr 2023 11:14:22 +0800
Subject: [PATCH] test: support TABBY_API_HOST in k6 tests

---
 Makefile                  |  4 ++--
 tests/default.loadtest.js | 17 ++++++++++-------
 tests/default.smoke.js    | 15 +++++++++++++--
 3 files changed, 25 insertions(+), 11 deletions(-)

diff --git a/Makefile b/Makefile
index bd1079b..b6132dc 100644
--- a/Makefile
+++ b/Makefile
@@ -19,8 +19,8 @@ $(PRE_COMMIT_HOOK):
 
 setup-development-environment: install-poetry $(PRE_COMMIT_HOOK)
 
-test-smoke:
+smoke:
 	k6 run tests/*.smoke.js
 
-test-loadtest:
+loadtest:
 	k6 run tests/*.loadtest.js
diff --git a/tests/default.loadtest.js b/tests/default.loadtest.js
index 2721205..3280376 100644
--- a/tests/default.loadtest.js
+++ b/tests/default.loadtest.js
@@ -1,14 +1,17 @@
-import http from 'k6/http';
-import { check, group, sleep } from 'k6';
+import http from "k6/http";
+import { check, group, sleep } from "k6";
 
 export const options = {
   stages: [
-    { duration: '5s', target: 10 }, // simulate ramp-up of traffic from 1 to 10 users over 30s.
-    { duration: '30s', target: 10 }, // stay at 10 users for 10 minutes
-    { duration: '5s', target: 0 }, // ramp-down to 0 users
+    { duration: "5s", target: 10 }, // simulate ramp-up of traffic from 1 to 10 users over 30s.
+    { duration: "30s", target: 10 }, // stay at 10 users for 10 minutes
+    { duration: "5s", target: 0 }, // ramp-down to 0 users
   ],
+  hosts: {
+    "api.tabbyml.com": __ENV.TABBY_API_HOST || "localhost:5000",
+  },
   thresholds: {
-    'http_req_duration': ['p(99)<1000'], // 99% of requests must complete below 1000ms
+    http_req_duration: ["p(99)<1000"], // 99% of requests must complete below 1000ms
   },
 };
 
@@ -17,7 +20,7 @@ export default () => {
     prompt: "def binarySearch(arr, left, right, x):\n    mid = (left +",
   });
   const headers = { "Content-Type": "application/json" };
-  const res = http.post("http://localhost:5000/v1/completions", payload, {
+  const res = http.post("http://api.tabbyml.com/v1/completions", payload, {
     headers,
   });
   check(res, { success: (r) => r.status === 200 });
diff --git a/tests/default.smoke.js b/tests/default.smoke.js
index d368951..30697d9 100644
--- a/tests/default.smoke.js
+++ b/tests/default.smoke.js
@@ -1,13 +1,24 @@
 import http from "k6/http";
-import { check } from "k6";
+import { check, sleep } from "k6";
+
+export const options = {
+    stages: [
+        {duration: '3s', target: 5},
+    ],
+    hosts: {
+        'api.tabbyml.com': __ENV.TABBY_API_HOST || "localhost:5000"
+    },
+};
+const SLEEP_DURATION = 1;
 
 export default function () {
   const payload = JSON.stringify({
     prompt: "def binarySearch(arr, left, right, x):\n    mid = (left +",
   });
   const headers = { "Content-Type": "application/json" };
-  const res = http.post("http://localhost:5000/v1/completions", payload, {
+  const res = http.post("http://api.tabbyml.com/v1/completions", payload, {
     headers,
   });
   check(res, { success: (r) => r.status === 200 });
+  sleep(SLEEP_DURATION)
 }