harness_zbt/backend/app/domains/attribution/repository.py

from typing import Protocol

from app.domains.attribution.models import (
    LeadAttribution,
    OrderAttribution,
    PromotionCode,
    Salesperson,
)


class AttributionRepository(Protocol):
    def save_salesperson(self, salesperson: Salesperson) -> None: ...

    def list_salespersons(self) -> list[Salesperson]: ...

    def save_promotion_code(self, promotion_code: PromotionCode) -> None: ...

    def get_promotion_code(self, code: str) -> PromotionCode | None: ...

    def list_promotion_codes(self) -> list[PromotionCode]: ...

    def save_lead(self, lead: LeadAttribution) -> None: ...

    def get_lead_by_user(self, h5_user_id: str) -> LeadAttribution | None: ...

    def list_leads(self) -> list[LeadAttribution]: ...

    def save_order_attribution(self, attribution: OrderAttribution) -> None: ...

    def get_order_attribution(self, order_id: str) -> OrderAttribution | None: ...

    def list_order_attributions(self) -> list[OrderAttribution]: ...


class InMemoryAttributionRepository:
    def __init__(self) -> None:
        self._salespersons: dict[str, Salesperson] = {}
        self._promotion_codes: dict[str, PromotionCode] = {}
        self._leads: dict[str, LeadAttribution] = {}
        self._orders: dict[str, OrderAttribution] = {}

    def save_salesperson(self, salesperson: Salesperson) -> None:
        self._salespersons[salesperson.id] = salesperson

    def list_salespersons(self) -> list[Salesperson]:
        return sorted(self._salespersons.values(), key=lambda item: item.code)

    def save_promotion_code(self, promotion_code: PromotionCode) -> None:
        self._promotion_codes[promotion_code.code.upper()] = promotion_code

    def get_promotion_code(self, code: str) -> PromotionCode | None:
        return self._promotion_codes.get(code.upper())

    def list_promotion_codes(self) -> list[PromotionCode]:
        return sorted(self._promotion_codes.values(), key=lambda item: item.code)

    def save_lead(self, lead: LeadAttribution) -> None:
        self._leads[lead.h5_user_id] = lead

    def get_lead_by_user(self, h5_user_id: str) -> LeadAttribution | None:
        return self._leads.get(h5_user_id)

    def list_leads(self) -> list[LeadAttribution]:
        return list(self._leads.values())

    def save_order_attribution(self, attribution: OrderAttribution) -> None:
        self._orders[attribution.order_id] = attribution

    def get_order_attribution(self, order_id: str) -> OrderAttribution | None:
        return self._orders.get(order_id)

    def list_order_attributions(self) -> list[OrderAttribution]:
        return list(self._orders.values())