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

from typing import Protocol

from app.domains.enrollment.models import (
    AsyncTask,
    EnrollmentDraft,
    EnrollmentOrder,
    OutboxEvent,
    PaymentTransaction,
    Policy,
    Quote,
    UserConfirmation,
)


class EnrollmentRepository(Protocol):
    def save_quote(self, quote: Quote) -> None: ...

    def get_quote(self, quote_id: str) -> Quote | None: ...

    def save_draft(self, draft: EnrollmentDraft) -> None: ...

    def get_draft(self, draft_id: str) -> EnrollmentDraft | None: ...

    def save_confirmation(self, confirmation: UserConfirmation) -> None: ...

    def get_confirmation_by_token_hash(self, token_hash: str) -> UserConfirmation | None: ...

    def save_order(self, order: EnrollmentOrder) -> None: ...

    def get_order(self, order_id: str) -> EnrollmentOrder | None: ...

    def find_order_by_idempotency(
        self, user_id: str, idempotency_key: str
    ) -> EnrollmentOrder | None: ...

    def list_orders(self, user_id: str | None = None) -> list[EnrollmentOrder]: ...

    def save_payment(self, payment: PaymentTransaction) -> None: ...

    def get_payment(self, payment_id: str) -> PaymentTransaction | None: ...

    def get_payment_by_no(self, payment_no: str) -> PaymentTransaction | None: ...

    def get_payment_by_order(self, order_id: str) -> PaymentTransaction | None: ...

    def find_payment_by_idempotency(
        self, order_id: str, idempotency_key: str
    ) -> PaymentTransaction | None: ...

    def save_policy(self, policy: Policy) -> None: ...

    def get_policy_by_order(self, order_id: str) -> Policy | None: ...

    def list_policies(self, user_id: str | None = None) -> list[Policy]: ...

    def save_outbox_event(self, event: OutboxEvent) -> None: ...

    def save_task(self, task: AsyncTask) -> None: ...

    def get_task_by_idempotency(self, idempotency_key: str) -> AsyncTask | None: ...


class InMemoryEnrollmentRepository:
    def __init__(self) -> None:
        self._quotes: dict[str, Quote] = {}
        self._drafts: dict[str, EnrollmentDraft] = {}
        self._confirmations: dict[str, UserConfirmation] = {}
        self._orders: dict[str, EnrollmentOrder] = {}
        self._payments: dict[str, PaymentTransaction] = {}
        self._policies: dict[str, Policy] = {}
        self._outbox_events: dict[str, OutboxEvent] = {}
        self._tasks: dict[str, AsyncTask] = {}

    def save_quote(self, quote: Quote) -> None:
        self._quotes[quote.id] = quote

    def get_quote(self, quote_id: str) -> Quote | None:
        return self._quotes.get(quote_id)

    def save_draft(self, draft: EnrollmentDraft) -> None:
        self._drafts[draft.id] = draft

    def get_draft(self, draft_id: str) -> EnrollmentDraft | None:
        return self._drafts.get(draft_id)

    def save_confirmation(self, confirmation: UserConfirmation) -> None:
        self._confirmations[confirmation.id] = confirmation

    def get_confirmation_by_token_hash(self, token_hash: str) -> UserConfirmation | None:
        return next(
            (
                confirmation
                for confirmation in self._confirmations.values()
                if confirmation.token_hash == token_hash
            ),
            None,
        )

    def save_order(self, order: EnrollmentOrder) -> None:
        self._orders[order.id] = order

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

    def find_order_by_idempotency(
        self, user_id: str, idempotency_key: str
    ) -> EnrollmentOrder | None:
        return next(
            (
                order
                for order in self._orders.values()
                if order.user_id == user_id and order.idempotency_key == idempotency_key
            ),
            None,
        )

    def list_orders(self, user_id: str | None = None) -> list[EnrollmentOrder]:
        return [
            order for order in self._orders.values() if user_id is None or order.user_id == user_id
        ]

    def save_payment(self, payment: PaymentTransaction) -> None:
        self._payments[payment.id] = payment

    def get_payment(self, payment_id: str) -> PaymentTransaction | None:
        return self._payments.get(payment_id)

    def get_payment_by_no(self, payment_no: str) -> PaymentTransaction | None:
        return next(
            (payment for payment in self._payments.values() if payment.payment_no == payment_no),
            None,
        )

    def get_payment_by_order(self, order_id: str) -> PaymentTransaction | None:
        return next(
            (payment for payment in self._payments.values() if payment.order_id == order_id),
            None,
        )

    def find_payment_by_idempotency(
        self, order_id: str, idempotency_key: str
    ) -> PaymentTransaction | None:
        return next(
            (
                payment
                for payment in self._payments.values()
                if payment.order_id == order_id and payment.idempotency_key == idempotency_key
            ),
            None,
        )

    def save_policy(self, policy: Policy) -> None:
        self._policies[policy.id] = policy

    def get_policy_by_order(self, order_id: str) -> Policy | None:
        return next(
            (policy for policy in self._policies.values() if policy.order_id == order_id),
            None,
        )

    def list_policies(self, user_id: str | None = None) -> list[Policy]:
        return [
            policy
            for policy in self._policies.values()
            if user_id is None or policy.user_id == user_id
        ]

    def save_outbox_event(self, event: OutboxEvent) -> None:
        self._outbox_events[event.id] = event

    def save_task(self, task: AsyncTask) -> None:
        self._tasks[task.id] = task

    def get_task_by_idempotency(self, idempotency_key: str) -> AsyncTask | None:
        return next(
            (task for task in self._tasks.values() if task.idempotency_key == idempotency_key),
            None,
        )