test_customer_segment.py

import pytest
from pytest_bdd import given, when, then, scenario, parsers, scenarios
from demo_project.transformers.user_segment import get_user_segment
import pandas as pd


@scenario(
    "test_customer_segment.feature",
    "Label customers as premium if they are between 19 and 60 and made high past purchases",
)
def test_label_customers_high_past_purchases():
    pass


@given(
    parsers.parse("a customer is {age:d} years old"),
    target_fixture="customer_age",
)
def customer_age(age):
    return age


@given(
    parsers.parse("the customer has {income_level} income"),
    target_fixture="customer_income",
)
def customer_income(income_level):
    return income_level


@given(
    "the customer has made high past purchases",
    target_fixture="customer_purchases",
)
def customer_purchases():
    return "high"


@when(
    "the segmentation process is executed",
    target_fixture="segmentation_execution",
)
def segmentation_execution(customer_age, customer_income, customer_purchases):
    df = pd.DataFrame(
        {
            "age": [customer_age],
            "income": [customer_income],
            "past_purchases": [customer_purchases],
        },
        columns=["age", "income", "past_purchases"],
    )
    return get_user_segment(df)


@then("the customer should be labeled as premium")
def verify_segmentation(segmentation_execution):
    assert segmentation_execution.iloc[0]["segment"] == "premium"