feat(loadtest): added loadtesting with k6
This commit is contained in:
@@ -0,0 +1,287 @@
|
||||
import json
|
||||
from decimal import Decimal
|
||||
from typing import override
|
||||
|
||||
from django.core.management.base import (
|
||||
BaseCommand,
|
||||
CommandError,
|
||||
CommandParser,
|
||||
)
|
||||
|
||||
from apps.experiments.models import Experiment, ExperimentStatus
|
||||
from apps.experiments.services import (
|
||||
experiment_approve,
|
||||
experiment_create,
|
||||
experiment_reopen,
|
||||
experiment_resume,
|
||||
experiment_start,
|
||||
experiment_submit_for_review,
|
||||
variant_create,
|
||||
)
|
||||
from apps.flags.models import FeatureFlag
|
||||
from apps.flags.selectors import feature_flag_get_by_key
|
||||
from apps.flags.services import feature_flag_create
|
||||
from apps.reviews.models import ApproverGroup
|
||||
from apps.users.models import User, UserRole
|
||||
|
||||
|
||||
class Command(BaseCommand):
|
||||
help = "Prepare deterministic fixture for k6 decide load tests"
|
||||
|
||||
@override
|
||||
def add_arguments(self, parser: CommandParser) -> None:
|
||||
parser.add_argument(
|
||||
"--run-id",
|
||||
required=True,
|
||||
type=str,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--owner",
|
||||
default="experimenter",
|
||||
type=str,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--approver",
|
||||
default="approver",
|
||||
type=str,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--flag-key",
|
||||
default=None,
|
||||
type=str,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--experiment-name",
|
||||
default=None,
|
||||
type=str,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--targeting-rules",
|
||||
default='country == "US"',
|
||||
type=str,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--json",
|
||||
action="store_true",
|
||||
default=False,
|
||||
)
|
||||
|
||||
@override
|
||||
def handle(self, *args, **options) -> None:
|
||||
run_id_raw: str = options["run_id"]
|
||||
owner_username: str = options["owner"]
|
||||
approver_username: str = options["approver"]
|
||||
flag_key_override: str | None = options["flag_key"]
|
||||
experiment_name_override: str | None = options["experiment_name"]
|
||||
targeting_rules: str = options["targeting_rules"]
|
||||
is_json: bool = options["json"]
|
||||
|
||||
run_id = self._normalize_run_id(run_id_raw)
|
||||
owner = self._load_user(owner_username, UserRole.EXPERIMENTER)
|
||||
approver = self._load_user(approver_username, UserRole.APPROVER)
|
||||
self._ensure_approver_group(owner, approver)
|
||||
|
||||
flag_key = flag_key_override or f"k6_{run_id}_flag"
|
||||
experiment_name = experiment_name_override or f"k6_{run_id}_experiment"
|
||||
|
||||
flag = self._ensure_flag(flag_key, run_id)
|
||||
experiment, created = self._resolve_experiment(
|
||||
flag=flag,
|
||||
owner=owner,
|
||||
name=experiment_name,
|
||||
targeting_rules=targeting_rules,
|
||||
)
|
||||
if created:
|
||||
self._ensure_variants(experiment=experiment, owner=owner)
|
||||
|
||||
experiment = self._ensure_running(
|
||||
experiment=experiment,
|
||||
owner=owner,
|
||||
approver=approver,
|
||||
)
|
||||
|
||||
payload = {
|
||||
"run_id": run_id,
|
||||
"flag_id": str(flag.pk),
|
||||
"flag_key": flag.key,
|
||||
"experiment_id": str(experiment.pk),
|
||||
"experiment_status": experiment.status,
|
||||
"owner": owner.username,
|
||||
"approver": approver.username,
|
||||
"subject_attributes": {"country": "US"},
|
||||
}
|
||||
|
||||
if is_json:
|
||||
self.stdout.write(json.dumps(payload))
|
||||
return
|
||||
|
||||
self.stdout.write(
|
||||
self.style.SUCCESS(
|
||||
f"k6 fixture ready: flag_key={flag.key}, "
|
||||
f"experiment_id={experiment.pk}, status={experiment.status}"
|
||||
)
|
||||
)
|
||||
self.stdout.write(json.dumps(payload, indent=2))
|
||||
|
||||
def _normalize_run_id(self, value: str) -> str:
|
||||
normalized = "".join(
|
||||
ch.lower() if ch.isalnum() else "_" for ch in value.strip()
|
||||
).strip("_")
|
||||
if not normalized:
|
||||
raise CommandError("run-id cannot be empty.")
|
||||
if not normalized[0].isalpha():
|
||||
normalized = f"r_{normalized}"
|
||||
return normalized
|
||||
|
||||
def _load_user(self, username: str, expected_role: str) -> User:
|
||||
user = User.objects.filter(username=username).first()
|
||||
if user is None:
|
||||
raise CommandError(
|
||||
f"User '{username}' was not found. Seed users before running."
|
||||
)
|
||||
if user.role != expected_role:
|
||||
raise CommandError(
|
||||
f"User '{username}' must have role '{expected_role}'."
|
||||
)
|
||||
return user
|
||||
|
||||
def _ensure_approver_group(self, owner: User, approver: User) -> None:
|
||||
group, _ = ApproverGroup.objects.get_or_create(
|
||||
experimenter=owner,
|
||||
defaults={"min_approvals": 1},
|
||||
)
|
||||
if group.min_approvals != 1:
|
||||
group.min_approvals = 1
|
||||
group.save(update_fields=["min_approvals", "updated_at"])
|
||||
if not group.approvers.filter(pk=approver.pk).exists():
|
||||
group.approvers.add(approver)
|
||||
|
||||
def _ensure_flag(self, key: str, run_id: str) -> FeatureFlag:
|
||||
flag = feature_flag_get_by_key(key)
|
||||
if flag:
|
||||
return flag
|
||||
return feature_flag_create(
|
||||
key=key,
|
||||
name=f"k6 {run_id} decide",
|
||||
value_type="string",
|
||||
default_value="control",
|
||||
)
|
||||
|
||||
def _resolve_experiment(
|
||||
self,
|
||||
*,
|
||||
flag: FeatureFlag,
|
||||
owner: User,
|
||||
name: str,
|
||||
targeting_rules: str,
|
||||
) -> tuple[Experiment, bool]:
|
||||
reusable = (
|
||||
Experiment.objects.filter(
|
||||
flag=flag,
|
||||
status__in=(
|
||||
ExperimentStatus.RUNNING,
|
||||
ExperimentStatus.PAUSED,
|
||||
ExperimentStatus.APPROVED,
|
||||
ExperimentStatus.IN_REVIEW,
|
||||
ExperimentStatus.DRAFT,
|
||||
ExperimentStatus.REJECTED,
|
||||
),
|
||||
)
|
||||
.order_by("-created_at")
|
||||
.first()
|
||||
)
|
||||
if reusable:
|
||||
return reusable, False
|
||||
|
||||
experiment = experiment_create(
|
||||
flag=flag,
|
||||
name=name,
|
||||
owner=owner,
|
||||
description="k6 decide benchmark fixture",
|
||||
hypothesis="k6 baseline",
|
||||
traffic_allocation=Decimal("100.00"),
|
||||
targeting_rules=targeting_rules,
|
||||
)
|
||||
return experiment, True
|
||||
|
||||
def _ensure_variants(self, *, experiment: Experiment, owner: User) -> None:
|
||||
if experiment.variants.exists():
|
||||
return
|
||||
variant_create(
|
||||
experiment=experiment,
|
||||
user=owner,
|
||||
name="control",
|
||||
value="control",
|
||||
weight=Decimal("50.00"),
|
||||
is_control=True,
|
||||
)
|
||||
variant_create(
|
||||
experiment=experiment,
|
||||
user=owner,
|
||||
name="treatment",
|
||||
value="treatment",
|
||||
weight=Decimal("50.00"),
|
||||
is_control=False,
|
||||
)
|
||||
|
||||
def _ensure_running(
|
||||
self,
|
||||
*,
|
||||
experiment: Experiment,
|
||||
owner: User,
|
||||
approver: User,
|
||||
) -> Experiment:
|
||||
current = Experiment.objects.select_related("flag", "owner").get(
|
||||
pk=experiment.pk
|
||||
)
|
||||
status = current.status
|
||||
|
||||
if status in {
|
||||
ExperimentStatus.COMPLETED,
|
||||
ExperimentStatus.ARCHIVED,
|
||||
}:
|
||||
raise CommandError(
|
||||
"Reusable experiment is completed/archived. Use a new run-id."
|
||||
)
|
||||
|
||||
if status == ExperimentStatus.REJECTED:
|
||||
current = experiment_reopen(experiment=current, user=owner)
|
||||
status = current.status
|
||||
|
||||
if status == ExperimentStatus.DRAFT:
|
||||
current = experiment_submit_for_review(
|
||||
experiment=current,
|
||||
user=owner,
|
||||
)
|
||||
status = current.status
|
||||
|
||||
if status == ExperimentStatus.IN_REVIEW:
|
||||
if not current.approvals.filter(approver=approver).exists():
|
||||
current = experiment_approve(
|
||||
experiment=current,
|
||||
approver=approver,
|
||||
comment="k6 fixture approval",
|
||||
)
|
||||
status = current.status
|
||||
if status == ExperimentStatus.IN_REVIEW:
|
||||
raise CommandError(
|
||||
"Experiment still in_review after approval. "
|
||||
"Check review policy for owner."
|
||||
)
|
||||
|
||||
if status == ExperimentStatus.APPROVED:
|
||||
current = experiment_start(experiment=current, user=owner)
|
||||
status = current.status
|
||||
|
||||
if status == ExperimentStatus.PAUSED:
|
||||
current = experiment_resume(experiment=current, user=owner)
|
||||
status = current.status
|
||||
|
||||
if status != ExperimentStatus.RUNNING:
|
||||
current = Experiment.objects.get(pk=current.pk)
|
||||
raise CommandError(
|
||||
"Failed to move experiment to running. "
|
||||
f"Current={current.status}"
|
||||
)
|
||||
|
||||
return current
|
||||
Reference in New Issue
Block a user