Initial commitv3.3.4

2 files changed, 608 insertions, 0 deletions

diff --git a/test_utils/models/__init__.py b/test_utils/models/__init__.py
new file mode 100644
index 0000000..e69de29
--- /dev/null
+++ b/test_utils/models/__init__.py
diff --git a/test_utils/models/conftest.py b/test_utils/models/conftest.py
new file mode 100644
index 0000000..ecfd82b
--- /dev/null
+++ b/test_utils/models/conftest.py
@@ -0,0 +1,608 @@
+from datetime import datetime, timezone, timedelta
+from decimal import Decimal
+from random import randint, choice as randchoice
+from typing import Callable, TYPE_CHECKING, Optional, List, Dict
+from uuid import uuid4
+
+import pytest
+from pydantic import AwareDatetime, PositiveInt
+
+from generalresearch.models import Source
+from generalresearch.models.thl.definitions import (
+    WALL_ALLOWED_STATUS_STATUS_CODE,
+    Status,
+)
+from test_utils.managers.conftest import (
+    product_manager,
+    user_manager,
+    wall_manager,
+    session_manager,
+    gr_um,
+    membership_manager,
+    team_manager,
+    business_manager,
+    business_address_manager,
+)
+from generalresearch.models.thl.survey.model import Survey, Buyer
+
+if TYPE_CHECKING:
+    from generalresearch.models.thl.userhealth import AuditLog, AuditLogLevel
+    from generalresearch.models.thl.payout import UserPayoutEvent
+    from generalresearch.models.gr.authentication import GRUser, GRToken
+    from generalresearch.models.gr.team import Team, Membership
+    from generalresearch.models.gr.business import (
+        Business,
+        BusinessAddress,
+        BusinessBankAccount,
+    )
+    from generalresearch.models.thl.user import User
+    from generalresearch.models.thl.product import Product
+    from generalresearch.models.thl.session import Session, Wall
+    from generalresearch.currency import USDCent
+    from generalresearch.models.thl.product import (
+        PayoutConfig,
+        PayoutTransformation,
+        PayoutTransformationPercentArgs,
+    )
+    from generalresearch.models.thl.user_iphistory import IPRecord
+    from generalresearch.models.thl.ipinfo import IPGeoname, IPInformation
+
+
+# === THL ===
+
+
+@pytest.fixture(scope="function")
+def user(request, product_manager, user_manager, thl_web_rr) -> "User":
+    product = getattr(request, "product", None)
+
+    if product is None:
+        product = product_manager.create_dummy()
+
+    u = user_manager.create_dummy(product_id=product.id)
+    u.prefetch_product(pg_config=thl_web_rr)
+
+    return u
+
+
+@pytest.fixture
+def user_with_wallet(
+    request, user_factory, product_user_wallet_yes: "Product"
+) -> "User":
+    # A user on a product with user wallet enabled, but they have no money
+    return user_factory(product=product_user_wallet_yes)
+
+
+@pytest.fixture
+def user_with_wallet_amt(request, user_factory, product_amt_true: "Product") -> "User":
+    # A user on a product with user wallet enabled, on AMT, but they have no money
+    return user_factory(product=product_amt_true)
+
+
+@pytest.fixture(scope="function")
+def user_factory(user_manager, thl_web_rr) -> Callable:
+    def _create_user(product: "Product", created: Optional[datetime] = None):
+        u = user_manager.create_dummy(product=product, created=created)
+        u.prefetch_product(pg_config=thl_web_rr)
+
+        return u
+
+    return _create_user
+
+
+@pytest.fixture(scope="function")
+def wall_factory(wall_manager) -> Callable:
+    def _create_wall(
+        session: "Session", wall_status: "Status", req_cpi: Optional[Decimal] = None
+    ):
+
+        assert session.started <= datetime.now(
+            tz=timezone.utc
+        ), "Session can't start in the future"
+
+        if session.wall_events:
+            # Subsequent Wall events
+            wall = session.wall_events[-1]
+            assert not wall.finished, "Can't add new Walls until prior finishes"
+            # wall_started = last_wall.started + timedelta(milliseconds=1)
+        else:
+            # First Wall Event in a session
+            wall_started = session.started + timedelta(milliseconds=1)
+
+            wall = wall_manager.create_dummy(
+                session_id=session.id,
+                user_id=session.user_id,
+                started=wall_started,
+                req_cpi=req_cpi,
+            )
+            session.append_wall_event(w=wall)
+
+        options = list(WALL_ALLOWED_STATUS_STATUS_CODE.get(wall_status, {}))
+        wall.finish(
+            finished=wall.started + timedelta(seconds=randint(a=60 * 2, b=60 * 10)),
+            status=wall_status,
+            status_code_1=randchoice(options),
+        )
+
+        return wall
+
+    return _create_wall
+
+
+@pytest.fixture(scope="function")
+def wall(session, user, wall_manager) -> Optional["Wall"]:
+    from generalresearch.models.thl.task_status import StatusCode1
+
+    wall = wall_manager.create_dummy(session_id=session.id, user_id=user.user_id)
+    # thl_session.append_wall_event(wall)
+    wall.finish(
+        finished=wall.started + timedelta(seconds=randint(a=60 * 2, b=60 * 10)),
+        status=Status.COMPLETE,
+        status_code_1=StatusCode1.COMPLETE,
+    )
+    return wall
+
+
+@pytest.fixture(scope="function")
+def session_factory(
+    wall_factory, session_manager, wall_manager, utc_hour_ago
+) -> Callable:
+    from generalresearch.models.thl.session import Source
+
+    def _create_session(
+        user: "User",
+        # Wall details
+        wall_count: int = 5,
+        wall_req_cpi: Decimal = Decimal(".50"),
+        wall_req_cpis: Optional[List[Decimal]] = None,
+        wall_statuses: Optional[List[Status]] = None,
+        wall_source: Source = Source.TESTING,
+        # Session details
+        final_status: Status = Status.COMPLETE,
+        started: datetime = utc_hour_ago,
+    ) -> "Session":
+        if wall_req_cpis:
+            assert len(wall_req_cpis) == wall_count
+        if wall_statuses:
+            assert len(wall_statuses) == wall_count
+
+        s = session_manager.create_dummy(started=started, user=user, country_iso="us")
+        for idx in range(wall_count):
+            if idx == 0:
+                # First Wall Event in a session
+                wall_started = s.started + timedelta(milliseconds=1)
+            else:
+                # Subsequent Wall events
+                last_wall = s.wall_events[-1]
+                assert last_wall.finished, "Can't add new Walls until prior finishes"
+                wall_started = last_wall.started + timedelta(milliseconds=1)
+
+            w = wall_manager.create_dummy(
+                session_id=s.id,
+                source=wall_source,
+                user_id=s.user_id,
+                started=wall_started,
+                req_cpi=wall_req_cpis[idx] if wall_req_cpis else wall_req_cpi,
+            )
+            s.append_wall_event(w=w)
+
+            # If it's the last wall in the session, respect the final_status
+            #   value for the Session
+            if wall_statuses:
+                _final_status = wall_statuses[idx]
+            else:
+                _final_status = final_status if idx == wall_count - 1 else Status.FAIL
+
+            options = list(WALL_ALLOWED_STATUS_STATUS_CODE.get(_final_status, {}))
+            wall_manager.finish(
+                wall=w,
+                status=_final_status,
+                status_code_1=randchoice(options),
+                finished=w.started + timedelta(seconds=randint(a=60 * 2, b=60 * 10)),
+            )
+
+        return s
+
+    return _create_session
+
+
+@pytest.fixture(scope="function")
+def finished_session_factory(
+    session_factory, session_manager, utc_hour_ago
+) -> Callable:
+    from generalresearch.models.thl.session import Source
+
+    def _create_finished_session(
+        user: "User",
+        # Wall details
+        wall_count: int = 5,
+        wall_req_cpi: Decimal = Decimal(".50"),
+        wall_req_cpis: Optional[List[Decimal]] = None,
+        wall_statuses: Optional[List[Status]] = None,
+        wall_source: Source = Source.TESTING,
+        # Session details
+        final_status: Status = Status.COMPLETE,
+        started: datetime = utc_hour_ago,
+    ) -> "Session":
+        s: Session = session_factory(
+            user=user,
+            wall_count=wall_count,
+            wall_req_cpi=wall_req_cpi,
+            wall_req_cpis=wall_req_cpis,
+            wall_statuses=wall_statuses,
+            wall_source=wall_source,
+            final_status=final_status,
+            started=started,
+        )
+        status, status_code_1 = s.determine_session_status()
+        thl_net, commission_amount, bp_pay, user_pay = s.determine_payments()
+        session_manager.finish_with_status(
+            s,
+            finished=s.wall_events[-1].finished,
+            payout=bp_pay,
+            user_payout=user_pay,
+            status=status,
+            status_code_1=status_code_1,
+        )
+        return s
+
+    return _create_finished_session
+
+
+@pytest.fixture(scope="function")
+def session(user, session_manager, wall_manager) -> "Session":
+    from generalresearch.models.thl.session import Wall, Session
+
+    session: Session = session_manager.create_dummy(user=user, country_iso="us")
+    wall: Wall = wall_manager.create_dummy(
+        session_id=session.id,
+        user_id=session.user_id,
+        started=session.started,
+    )
+    session.append_wall_event(w=wall)
+
+    return session
+
+
+@pytest.fixture
+def product(request, product_manager) -> "Product":
+    from generalresearch.managers.thl.product import ProductManager
+
+    team = getattr(request, "team", None)
+    business = getattr(request, "business", None)
+
+    product_manager: ProductManager
+    return product_manager.create_dummy(
+        team_id=team.uuid if team else None,
+        business_id=business.uuid if business else None,
+    )
+
+
+@pytest.fixture
+def product_factory(product_manager) -> Callable:
+    def _create_product(
+        team: Optional["Team"] = None,
+        business: Optional["Business"] = None,
+        commission_pct: Decimal = Decimal("0.05"),
+    ):
+        return product_manager.create_dummy(
+            team_id=team.uuid if team else None,
+            business_id=business.uuid if business else None,
+            commission_pct=commission_pct,
+        )
+
+    return _create_product
+
+
+@pytest.fixture(scope="function")
+def payout_config(request) -> "PayoutConfig":
+    from generalresearch.models.thl.product import (
+        PayoutConfig,
+        PayoutTransformation,
+        PayoutTransformationPercentArgs,
+    )
+
+    return (
+        request.param
+        if hasattr(request, "payout_config")
+        else PayoutConfig(
+            payout_format="${payout/100:.2f}",
+            payout_transformation=PayoutTransformation(
+                f="payout_transformation_percent",
+                kwargs=PayoutTransformationPercentArgs(pct=0.40),
+            ),
+        )
+    )
+
+
+@pytest.fixture(scope="function")
+def product_user_wallet_yes(payout_config, product_manager) -> "Product":
+    from generalresearch.models.thl.product import UserWalletConfig
+    from generalresearch.managers.thl.product import ProductManager
+
+    product_manager: ProductManager
+    return product_manager.create_dummy(
+        payout_config=payout_config, user_wallet_config=UserWalletConfig(enabled=True)
+    )
+
+
+@pytest.fixture(scope="function")
+def product_user_wallet_no(product_manager) -> "Product":
+    from generalresearch.models.thl.product import UserWalletConfig
+    from generalresearch.managers.thl.product import ProductManager
+
+    product_manager: ProductManager
+    return product_manager.create_dummy(
+        user_wallet_config=UserWalletConfig(enabled=False)
+    )
+
+
+@pytest.fixture(scope="function")
+def product_amt_true(product_manager, payout_config) -> "Product":
+    from generalresearch.models.thl.product import UserWalletConfig
+
+    return product_manager.create_dummy(
+        user_wallet_config=UserWalletConfig(amt=True, enabled=True),
+        payout_config=payout_config,
+    )
+
+
+@pytest.fixture(scope="function")
+def bp_payout_factory(
+    thl_lm, product_manager, business_payout_event_manager
+) -> Callable:
+    def _create_bp_payout(
+        product: Optional["Product"] = None,
+        amount: Optional["USDCent"] = None,
+        ext_ref_id: Optional[str] = None,
+        created: Optional[AwareDatetime] = None,
+        skip_wallet_balance_check: bool = False,
+        skip_one_per_day_check: bool = False,
+    ) -> "UserPayoutEvent":
+        from generalresearch.currency import USDCent
+
+        product = product or product_manager.create_dummy()
+        amount = amount or USDCent(randint(1, 99_99))
+
+        return business_payout_event_manager.create_bp_payout_event(
+            thl_ledger_manager=thl_lm,
+            product=product,
+            amount=amount,
+            ext_ref_id=ext_ref_id,
+            created=created,
+            skip_wallet_balance_check=skip_wallet_balance_check,
+            skip_one_per_day_check=skip_one_per_day_check,
+        )
+
+    return _create_bp_payout
+
+
+# === GR ===
+
+
+@pytest.fixture(scope="function")
+def business(request, business_manager) -> "Business":
+    from generalresearch.managers.gr.business import BusinessManager
+
+    business_manager: BusinessManager
+    return business_manager.create_dummy()
+
+
+@pytest.fixture(scope="function")
+def business_address(request, business, business_address_manager) -> "BusinessAddress":
+    from generalresearch.managers.gr.business import BusinessAddressManager
+
+    business_address_manager: BusinessAddressManager
+    return business_address_manager.create_dummy(business_id=business.id)
+
+
+@pytest.fixture(scope="function")
+def business_bank_account(
+    request, business, business_bank_account_manager
+) -> "BusinessBankAccount":
+    from generalresearch.managers.gr.business import BusinessBankAccountManager
+
+    business_bank_account_manager: BusinessBankAccountManager
+    return business_bank_account_manager.create_dummy(business_id=business.id)
+
+
+@pytest.fixture(scope="function")
+def team(request, team_manager) -> "Team":
+    from generalresearch.managers.gr.team import TeamManager
+
+    team_manager: TeamManager
+    return team_manager.create_dummy()
+
+
+@pytest.fixture(scope="function")
+def gr_user(gr_um) -> "GRUser":
+    from generalresearch.managers.gr.authentication import GRUserManager
+
+    gr_um: GRUserManager
+    return gr_um.create_dummy()
+
+
+@pytest.fixture(scope="function")
+def gr_user_cache(gr_user, gr_db, thl_web_rr, gr_redis_config):
+    gr_user.set_cache(
+        pg_config=gr_db, thl_web_rr=thl_web_rr, redis_config=gr_redis_config
+    )
+    return gr_user
+
+
+@pytest.fixture(scope="function")
+def gr_user_factory(gr_um) -> Callable:
+    def _create_gr_user():
+        return gr_um.create_dummy()
+
+    return _create_gr_user
+
+
+@pytest.fixture()
+def gr_user_token(gr_user, gr_tm, gr_db) -> "GRToken":
+    gr_tm.create(user_id=gr_user.id)
+    gr_user.prefetch_token(pg_config=gr_db)
+
+    return gr_user.token
+
+
+@pytest.fixture()
+def gr_user_token_header(gr_user_token) -> Dict:
+    return gr_user_token.auth_header
+
+
+@pytest.fixture(scope="function")
+def membership(request, team, gr_user, team_manager) -> "Membership":
+    assert team.id, "Team must be saved"
+    assert gr_user.id, "GRUser must be saved"
+    return team_manager.add_user(team=team, gr_user=gr_user)
+
+
+@pytest.fixture(scope="function")
+def membership_factory(
+    team: "Team", gr_user: "GRUser", membership_manager, team_manager, gr_um
+) -> Callable:
+    from generalresearch.managers.gr.team import MembershipManager
+
+    membership_manager: MembershipManager
+
+    def _create_membership(**kwargs):
+        _team = kwargs.get("team", team_manager.create_dummy())
+        _gr_user = kwargs.get("gr_user", gr_um.create_dummy())
+
+        return membership_manager.create(team=_team, gr_user=_gr_user)
+
+    return _create_membership
+
+
+@pytest.fixture(scope="function")
+def audit_log(audit_log_manager, user) -> "AuditLog":
+    from generalresearch.managers.thl.userhealth import AuditLogManager
+
+    audit_log_manager: AuditLogManager
+    return audit_log_manager.create_dummy(user_id=user.user_id)
+
+
+@pytest.fixture(scope="function")
+def audit_log_factory(audit_log_manager) -> Callable:
+    from generalresearch.managers.thl.userhealth import AuditLogManager
+
+    audit_log_manager: AuditLogManager
+
+    def _create_audit_log(
+        user_id: PositiveInt,
+        level: Optional["AuditLogLevel"] = None,
+        event_type: Optional[str] = None,
+        event_msg: Optional[str] = None,
+        event_value: Optional[float] = None,
+    ):
+        return audit_log_manager.create_dummy(
+            user_id=user_id,
+            level=level,
+            event_type=event_type,
+            event_msg=event_msg,
+            event_value=event_value,
+        )
+
+    return _create_audit_log
+
+
+@pytest.fixture(scope="function")
+def ip_geoname(ip_geoname_manager) -> "IPGeoname":
+    from generalresearch.managers.thl.ipinfo import IPGeonameManager
+
+    ip_geoname_manager: IPGeonameManager
+    return ip_geoname_manager.create_dummy()
+
+
+@pytest.fixture(scope="function")
+def ip_information(ip_information_manager, ip_geoname) -> "IPInformation":
+    from generalresearch.managers.thl.ipinfo import IPInformationManager
+
+    ip_information_manager: IPInformationManager
+    return ip_information_manager.create_dummy(
+        geoname_id=ip_geoname.geoname_id, country_iso=ip_geoname.country_iso
+    )
+
+
+@pytest.fixture(scope="function")
+def ip_information_factory(ip_information_manager) -> Callable:
+    from generalresearch.managers.thl.ipinfo import IPInformationManager
+
+    ip_information_manager: IPInformationManager
+
+    def _create_ip_info(ip: str, geoname: "IPGeoname", **kwargs):
+        return ip_information_manager.create_dummy(
+            ip=ip,
+            geoname_id=geoname.geoname_id,
+            country_iso=geoname.country_iso,
+            **kwargs,
+        )
+
+    return _create_ip_info
+
+
+@pytest.fixture(scope="function")
+def ip_record(ip_record_manager, ip_geoname, user) -> "IPRecord":
+    from generalresearch.managers.thl.userhealth import IPRecordManager
+
+    ip_record_manager: IPRecordManager
+
+    return ip_record_manager.create_dummy(user_id=user.user_id)
+
+
+@pytest.fixture(scope="function")
+def ip_record_factory(ip_record_manager, user) -> Callable:
+    from generalresearch.managers.thl.userhealth import IPRecordManager
+
+    ip_record_manager: IPRecordManager
+
+    def _create_ip_record(user_id: PositiveInt, ip: Optional[str] = None):
+        return ip_record_manager.create_dummy(user_id=user_id, ip=ip)
+
+    return _create_ip_record
+
+
+@pytest.fixture(scope="session")
+def buyer(buyer_manager) -> Buyer:
+    buyer_code = uuid4().hex
+    buyer_manager.bulk_get_or_create(source=Source.TESTING, codes=[buyer_code])
+    b = Buyer(
+        source=Source.TESTING, code=buyer_code, label=f"test-buyer-{buyer_code[:8]}"
+    )
+    buyer_manager.update(b)
+    return b
+
+
+@pytest.fixture(scope="session")
+def buyer_factory(buyer_manager) -> Callable:
+
+    def inner():
+        return buyer_manager.bulk_get_or_create(
+            source=Source.TESTING, codes=[uuid4().hex]
+        )[0]
+
+    return inner
+
+
+@pytest.fixture(scope="session")
+def survey(survey_manager, buyer) -> Survey:
+    s = Survey(source=Source.TESTING, survey_id=uuid4().hex, buyer_code=buyer.code)
+    survey_manager.create_bulk([s])
+    return s
+
+
+@pytest.fixture(scope="session")
+def survey_factory(survey_manager, buyer_factory) -> Callable:
+
+    def inner(buyer: Optional[Buyer] = None) -> Survey:
+        buyer = buyer or buyer_factory()
+        s = Survey(
+            source=Source.TESTING,
+            survey_id=uuid4().hex,
+            buyer_code=buyer.code,
+            buyer_id=buyer.id,
+        )
+        survey_manager.create_bulk([s])
+        return s
+
+    return inner