cvat/tests/perf/perfkit/k6_profile.py

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")