from locust import HttpUser, between, task
import random

EXAMPLE_PAYLOADS = [
    {
        "merchant_name":"TestA",
        "avg_amount_per_transaction":123.45,
        "day_of_week":2,
        "amount_deviation_from_location_mean":0,
        "transaction_category":"retail",
        "customer_no_transactions":0,
        "customer_lat":None,
        "transaction_type":"online",
        "customer_place_name":None,
        "merchant_id":None,
        "location":"New York",
        "customer_job":None,
        "age":None,
        "merchant_long":None,
        "amount_per_city_pop":0,
        "customer_long":None,
        "distance_customer_merchant":0,
        "transactions_per_customer_ratio":0,
        "customer_city_population":0,
        "merchant_lat":None,
        "customer_no_payments":0,
        "customer_no_orders":0,
        "payments_per_order_ratio":0,
        "hour_of_day":12,
        "amount":123.45,
        "customer_zip_code":None,
        "mean_amount_by_location":0,
        "fraud_rate_by_location":0,
        "customer_gender":None
    },
]

class PredictUser(HttpUser):
    wait_time = between(1, 3)

    @task(10)
    def predict(self):
        payload = random.choice(EXAMPLE_PAYLOADS)
        self.client.post("/predict", json=payload, name="POST /predict")

    @task(1)
    def get_root(self):
        self.client.get("/", name="GET /")