diff --git a/posthog/management/commands/generate_experiment_data.py b/posthog/management/commands/generate_experiment_data.py index 58601a74daaf7..291e78ad6ad58 100644 --- a/posthog/management/commands/generate_experiment_data.py +++ b/posthog/management/commands/generate_experiment_data.py @@ -2,6 +2,7 @@ import logging import random import time +from typing import Literal import uuid import json @@ -29,65 +30,114 @@ class ExperimentConfig(BaseModel): variants: dict[str, VariantConfig] -def get_default_experiment_config() -> ExperimentConfig: +def get_default_funnel_experiment_config() -> ExperimentConfig: return ExperimentConfig( - number_of_users=1000, + number_of_users=2000, start_timestamp=datetime.now() - timedelta(days=7), end_timestamp=datetime.now(), variants={ "control": VariantConfig( weight=0.5, - actions=[ActionConfig(event="$pageview", count=1, probability=0.75)], + actions=[ + ActionConfig(event="signup started", count=1, probability=1), + ActionConfig(event="signup completed", count=1, probability=0.25), + ], ), "test": VariantConfig( weight=0.5, - actions=[ActionConfig(event="$pageview", count=1, probability=1)], + actions=[ + ActionConfig(event="signup started", count=1, probability=1), + ActionConfig(event="signup completed", count=1, probability=0.35), + ], ), }, ) +def get_default_trends_experiment_config() -> ExperimentConfig: + return ExperimentConfig( + number_of_users=2000, + start_timestamp=datetime.now() - timedelta(days=7), + end_timestamp=datetime.now(), + variants={ + "control": VariantConfig( + weight=0.5, + actions=[ActionConfig(event="$pageview", count=5, probability=0.25)], + ), + "test": VariantConfig( + weight=0.5, + actions=[ActionConfig(event="$pageview", count=5, probability=0.35)], + ), + }, + ) + + +def get_default_config(type: Literal["funnel", "trends"]) -> ExperimentConfig: + match type: + case "funnel": + return get_default_funnel_experiment_config() + case "trends": + return get_default_trends_experiment_config() + case _: + raise ValueError(f"Invalid experiment type: {type}") + + class Command(BaseCommand): help = "Generate experiment test data" def add_arguments(self, parser): group = parser.add_mutually_exclusive_group(required=False) group.add_argument( - "--init-config", type=str, help="Initialize a new experiment configuration file at the specified path" + "--init-config", + type=str, + help="Initialize a new experiment configuration file at the specified path. Does not generate data.", + ) + group.add_argument( + "--type", + type=str, + choices=["trends", "funnel"], + default="trends", + help="Type of experiment data to generate or configuration to initialize.", ) experiment_group = parser.add_argument_group("experiment arguments") experiment_group.add_argument("--experiment-id", type=str, help="Experiment ID (feature flag name)") experiment_group.add_argument("--config", type=str, help="Path to experiment config file") - experiment_group.add_argument( - "--seed", type=str, required=False, help="Simulation seed for deterministic output" - ) def handle(self, *args, **options): # Make sure this runs in development environment only if not settings.DEBUG: raise ValueError("This command should only be run in development! DEBUG must be True.") - if config_path := options.get("init_config"): - with open(config_path, "w") as f: - f.write(get_default_experiment_config().model_dump_json(indent=2)) - logging.info(f"Created example configuration file at: {config_path}") + experiment_type = options.get("type") + + if init_config_path := options.get("init_config"): + with open(init_config_path, "w") as f: + f.write(get_default_config(experiment_type).model_dump_json(indent=2)) + logging.info(f"Created example {experiment_type} configuration file at: {init_config_path}") return experiment_id = options.get("experiment_id") config_path = options.get("config") - if not experiment_id or not config_path: - raise ValueError("Both --experiment-id and --config are required when not using --init-config") + # Validate required arguments + if not experiment_id: + raise ValueError("--experiment-id is missing!") + + if config_path is None and experiment_type is None: + raise ValueError("--config or --type trends|funnel is missing!") - with open(config_path) as config_file: - config_data = json.load(config_file) + if config_path: + with open(config_path) as config_file: + config_data = json.load(config_file) - try: - # Use the ExperimentConfig model to parse and validate the JSON data - experiment_config = ExperimentConfig(**config_data) - except ValidationError as e: - raise ValueError(f"Invalid configuration: {e}") + try: + # Use the ExperimentConfig model to parse and validate the JSON data + experiment_config = ExperimentConfig(**config_data) + except ValidationError as e: + raise ValueError(f"Invalid configuration: {e}") + else: + experiment_config = get_default_config(experiment_type) variants = list(experiment_config.variants.keys()) variant_counts = {variant: 0 for variant in variants}