WIP

posthog/management/commands/generate_experiment_data.py

@@ -2,6 +2,7 @@
 import logging
 import random
 import secrets
+import time
 import uuid
 
 from django.conf import settings
@@ -25,15 +26,18 @@ def handle(self, *args, **options):
             raise ValueError("This command should only be run in development! DEBUG must be True.")
 
         experiment_id = options.get("experiment_id")
+
+        # TODO: actually implement a seed
         seed = options.get("seed") or secrets.token_hex(16)
 
         if not experiment_id:
             raise ValueError("Experiment ID is required")
 
+        # TODO: this can be a config file taken as an argument
         experiment_config = {
             "experiment_id": experiment_id,
             "seed": seed,
-            "number_of_users": 10,
+            "number_of_users": 1000,
             "start_timestamp": datetime.now() - timedelta(days=7),
             "end_timestamp": datetime.now(),
             "variants": {
@@ -52,31 +56,39 @@ def handle(self, *args, **options):
             },
         }
 
+        variants = list(experiment_config["variants"].keys())
+        variant_counts = {variant: 0 for variant in variants}
         for _ in range(experiment_config["number_of_users"]):
             variant = random.choices(
-                list(experiment_config["variants"].keys()),
+                variants,
                 weights=[v["weight"] for v in experiment_config["variants"].values()],
             )[0]
+            variant_counts[variant] += 1
             distinct_id = uuid.uuid4()
             random_timestamp = random.uniform(
                 experiment_config["start_timestamp"], experiment_config["end_timestamp"] - timedelta(hours=1)
             )
-            print(f"Generating data for user {distinct_id} at {random_timestamp} with variant {variant}")
             posthoganalytics.capture(
                 distinct_id=distinct_id,
                 event="$feature_flag_called",
                 timestamp=random_timestamp,
                 properties={
-                    "feature_flag": experiment_config["experiment_id"],
-                    f"feature/{experiment_config['experiment_id']}": variant,
+                    "$feature_flag": experiment_config["experiment_id"],
+                    f"$feature/{experiment_config['experiment_id']}": variant,
                 },
             )
 
             for action in experiment_config["variants"][variant]["actions"]:
                 if random.random() < action["probability"]:
-                    print(f"Generating data for user {distinct_id} at {random_timestamp + timedelta(minutes=1)} with event {action['event']}")
                     posthoganalytics.capture(
                         distinct_id=distinct_id,
                         event=action["event"],
                         timestamp=random_timestamp + timedelta(minutes=1),
                     )
+
+        logging.info(f"Generated data for {experiment_config['experiment_id']} with seed {seed}")
+        logging.info(f"Variant counts: {variant_counts}")
+
+        # TODO: need to figure out how to wait for the data to be flushed. shutdown() doesn't work as expected.
+        time.sleep(10)
+        posthoganalytics.shutdown()