64 lines
1.9 KiB
Python
64 lines
1.9 KiB
Python
|
# Copyright (C) CVAT.ai Corporation
|
||
|
|
#
|
||
|
|
# SPDX-License-Identifier: MIT
|
||
|
|
from pathlib import Path
|
||
|
|
from typing import List, Optional, Union
|
||
|
|
|
||
|
|
from perfkit.config import PROMETHEUS_URL
|
||
|
|
|
||
|
|
_DEFAULT_TREND_STATS = "count,sum,min,max,avg,med,p(90),p(95),p(99)"
|
||
|
|
|
||
|
|
|
||
|
|
class K6Profile:
|
||
|
|
def __init__(self, test_file: Union[str, Path], enable_dashboard: bool = False) -> None:
|
||
|
|
self.test_file = Path(test_file)
|
||
|
|
self.env_vars: dict[str, str] = {
|
||
|
|
"K6_PROMETHEUS_RW_SERVER_URL": f"{PROMETHEUS_URL}/api/v1/write",
|
||
|
|
"K6_PROMETHEUS_RW_TREND_STATS": _DEFAULT_TREND_STATS,
|
||
|
|
}
|
||
|
|
if enable_dashboard:
|
||
|
|
self.env_vars["K6_WEB_DASHBOARD"] = "true"
|
||
|
|
self.args: List[str] = []
|
||
|
|
|
||
|
|
def add_arg(self, arg: str, value: Optional[str] = None) -> None:
|
||
|
|
self.args.append(arg)
|
||
|
|
if value is not None:
|
||
|
|
self.args.append(value)
|
||
|
|
|
||
|
|
def remove_arg(self, arg: str) -> None:
|
||
|
|
self.args.remove(arg)
|
||
|
|
|
||
|
|
def set_env(self, key: str, value: str) -> None:
|
||
|
|
self.env_vars[key] = value
|
||
|
|
|
||
|
|
def remove_env(self, key: str) -> None:
|
||
|
|
self.env_vars.pop(key, None)
|
||
|
|
|
||
|
|
def _build_env_list(self) -> List[str]:
|
||
|
|
result = []
|
||
|
|
for key, val in self.env_vars.items():
|
||
|
|
result.extend(["-e", f"{key}={val}"])
|
||
|
|
return result
|
||
|
|
|
||
|
|
def build_run_cmd(self) -> List[str]:
|
||
|
|
return [
|
||
|
|
"run",
|
||
|
|
*self._build_env_list(),
|
||
|
|
"--out",
|
||
|
|
"experimental-prometheus-rw",
|
||
|
|
"--address",
|
||
|
|
"0.0.0.0:6565",
|
||
|
|
"--tag",
|
||
|
|
f"testid={self.test_file.name}",
|
||
|
|
"--summary-export=/output/summary.json",
|
||
|
|
*self.args,
|
||
|
|
str(self.test_file),
|
||
|
|
]
|
||
|
|
|
||
|
|
def __repr__(self):
|
||
|
|
return f"<K6Config file={self.test_file.name} args={self.args} env={list(self.env_vars.keys())}>"
|
||
|
|
|
||
|
|
|
||
|
|
warmup_profile = K6Profile("tests/warmup.js")
|
||
|
|
tasks_regression_profile = K6Profile("tests/regression/tasks.js")
|