From 7958cf2e8771a0ecd7074661e84a803fba742725 Mon Sep 17 00:00:00 2001 From: stuppie Date: Mon, 5 Oct 2026 16:16:41 -0600 Subject: add some helpers so we dont break existing tests. Add a test for the bulk load with both kinds of products. --- generalresearch/models/thl/product.py | 151 ++++++++++------------------------ tests/models/thl/test_product.py | 141 ++++++++++++++++++++++++------- 2 files changed, 156 insertions(+), 136 deletions(-) diff --git a/generalresearch/models/thl/product.py b/generalresearch/models/thl/product.py index 7818e52..0330207 100644 --- a/generalresearch/models/thl/product.py +++ b/generalresearch/models/thl/product.py @@ -1088,7 +1088,6 @@ class Product(BaseModel, validate_assignment=True): self.bp_account = account # --- Prebuild --- - @staticmethod def get_pop_ledger_df( product_ids: Collection[UUIDStr], @@ -1158,10 +1157,26 @@ class Product(BaseModel, validate_assignment=True): ] return df - def prebuild_balance( + def prebuild_balance_individual( self, thl_lm: ThlLedgerManager, - pop_ledger_df: pd.DataFrame, + client: Client, + ds: GRLDatasets | None = None, + pop_ledger: PopLedgerMerge | None = None, + ) -> None: + if pop_ledger is None: + assert ds is not None + from generalresearch.incite.defaults import pop_ledger as plm + + pop_ledger = plm(ds=ds) + + pop_ledger_df = self.get_pop_ledger_df( + product_ids=[self.uuid], thl_lm=thl_lm, client=client, pop_ledger=pop_ledger + ) + self.prebuild_balance(thl_lm=thl_lm, pop_ledger_df=pop_ledger_df) + + def prebuild_balance( + self, thl_lm: ThlLedgerManager, pop_ledger_df: pd.DataFrame ) -> None: """ This returns the Product's Balances that are calculated across @@ -1185,7 +1200,7 @@ class Product(BaseModel, validate_assignment=True): absolutely required, Smart Retainer accounts are supported for any volume levels. """ - LOG.debug(f"Product.prebuild_balance_from_pop_ledger_df({self.uuid=})") + LOG.debug(f"Product.prebuild_balance({self.uuid=})") self.balance = None if self.bp_account is None: @@ -1322,7 +1337,6 @@ class Product(BaseModel, validate_assignment=True): def prebuild_payouts( self, - thl_lm: ThlLedgerManager, bp_pem: BrokerageProductPayoutEventManager, ) -> None: LOG.debug(f"Product.prebuild_payouts({self.uuid=})") @@ -1341,140 +1355,63 @@ class Product(BaseModel, validate_assignment=True): self.payouts_total = USDCent(sum([po.amount for po in self.payouts])) self.payouts_total_str = self.payouts_total.to_usd_str() - # def prebuild_pop(self): - # account = LM.get_account(qualified_name=f"{LM.currency.value}:bp_wallet:{product.id}") - # - # from main import data - # - # gv: GlobalVar = data["gv"] - # - # ddf = gv.pop_ledger.ddf( - # force_rr_latest=False, - # include_partial=True, - # columns=numerical_col_names + ["time_idx"], - # filters=[ - # ("account_id", "==", account.uuid), - # ("time_idx", ">=", rr.start), - # ], - # ) - # - # df = gv.dask_client.compute(collections=ddf, sync=True) - # df = df.set_index("time_idx").resample(rr.freq).sum() - # - # res = [] - # for index, row in df.iterrows(): - # index: pd.Timestamp - # row: pd.DataFrame - # - # dt = index.to_pydatetime().replace(tzinfo=None) - # instance = ProductBalances.from_pandas(row) - # - # res.append( - # { - # "time": dt, - # "payout": instance.payout / 100, - # "adjustment": instance.adjustment / 100, - # "expense": instance.expense / 100, - # "net": (instance.payout + instance.adjustment + instance.expense) / 100, - # } - # ) - # - # df = pd.DataFrame.from_records(res) - - # def financial( - # product: Product = Depends(product_from_path), - # rr: ReportRequest = Depends(rr_from_query), - # ) -> Any: - # account = LM.get_account(qualified_name=f"{LM.currency.value}:bp_wallet:{product.id}") - # - # from main import data - # - # gv: GlobalVar = data["gv"] - # - # ddf = gv.pop_ledger.ddf( - # force_rr_latest=False, - # include_partial=True, - # columns=numerical_col_names + ["time_idx", "account_id"], - # filters=[("account_id", "==", account.uuid), ("time_idx", ">=", rr.start)], - # ) - # - # df = gv.dask_client.compute(collections=ddf, sync=True) - # - # # We only do it this way so it's consistent with the Business.financial view - # df = df.groupby([pd.Grouper(key="time_idx", freq=rr.interval), "account_id"]).sum() - # return POPFinancial.list_from_pandas(df, accounts=[account]) - - # def payments(self): - # """Payments are the amount of money that General Research has sent - # the owner of this Product. - # - # These are typically ACH or Wire payments to company bank accounts. - # These are not respondent payments for Products where - # - # This is Provided in a standard list without any POP Grouping to show - # the exact time and amount of any Issued Payments. - # """ - # - # account = LM.get_account(qualified_name=f"{LM.currency.value}:bp_wallet:{product.id}") - # - # from main import data - # - # gv: GlobalVar = data["gv"] - # ddf = gv.pop_ledger.ddf( - # force_rr_latest=False, - # include_partial=True, - # columns=numerical_col_names + ["time_idx", "account_id"], - # filters=[("account_id", "==", account.uuid)], - # ) - # - # df = gv.dask_client.compute(collections=ddf, sync=True) - - # --- Methods --- def set_cache( self, thl_lm: ThlLedgerManager, - ds: GRLDatasets, client: Client, bp_pem: BrokerageProductPayoutEventManager, redis_config: RedisConfig, + ds: GRLDatasets | None = None, pop_ledger: PopLedgerMerge | None = None, - pop_ledger_df: pd.DataFrame | None = None, - ) -> None: + ): LOG.debug(f"Product.set_cache({self.uuid=})") if pop_ledger is None: + assert ds is not None from generalresearch.incite.defaults import pop_ledger as plm pop_ledger = plm(ds=ds) + pop_ledger_df = Product.get_pop_ledger_df( + product_ids=[self.uuid], + client=client, + thl_lm=thl_lm, + pop_ledger=pop_ledger, + ) + self.set_cache_from_pop_ledger_df( + thl_lm=thl_lm, + bp_pem=bp_pem, + redis_config=redis_config, + pop_ledger_df=pop_ledger_df, + ) + + def set_cache_from_pop_ledger_df( + self, + thl_lm: ThlLedgerManager, + bp_pem: BrokerageProductPayoutEventManager, + redis_config: RedisConfig, + pop_ledger_df: pd.DataFrame, + ) -> None: + LOG.debug(f"Product.set_cache({self.uuid=})") if self.bp_account is None: self.prefetch_bp_account(thl_lm=thl_lm) self.prebuild_balance( thl_lm=thl_lm, - client=client, - pop_ledger=pop_ledger, pop_ledger_df=pop_ledger_df, ) if self.balance: self.prebuild_private_balance( thl_lm=thl_lm, - client=client, - pop_ledger=pop_ledger, pop_ledger_df=pop_ledger_df, ) if self.user_wallet_enabled: self.prebuild_user_wallet_balances( - client=client, - pop_ledger=pop_ledger, pop_ledger_df=pop_ledger_df, ) - self.prebuild_payouts(thl_lm=thl_lm, bp_pem=bp_pem) + self.prebuild_payouts(bp_pem=bp_pem) self.prebuild_pop_financial( thl_lm=thl_lm, - ds=ds, - client=client, - pop_ledger=pop_ledger, pop_ledger_df=pop_ledger_df, ) diff --git a/tests/models/thl/test_product.py b/tests/models/thl/test_product.py index e45bd42..ce3242e 100644 --- a/tests/models/thl/test_product.py +++ b/tests/models/thl/test_product.py @@ -673,17 +673,17 @@ class TestProductFinancials: ) with pytest.raises(expected_exception=AssertionError) as cm: - p1.prebuild_balance( + p1.prebuild_balance_individual( thl_lm=thl_ledger_manager, ds=mnt_filepath, client=client_no_amm, ) - assert "Cannot build Product Balance" in str(cm.value) + assert "Cannot load Product POP ledger" in str(cm.value) ledger_collection.initial_load(client=None, sync=True) pop_ledger_merge.build(client=client_no_amm, ledger_coll=ledger_collection) - p1.prebuild_balance( + p1.prebuild_balance_individual( thl_lm=thl_ledger_manager, ds=mnt_filepath, client=client_no_amm, @@ -700,7 +700,6 @@ class TestProductFinancials: body, content_type = p1.balance.to_prometheus() p1.prebuild_payouts( - thl_lm=thl_ledger_manager, bp_pem=brokerage_product_payout_event_manager, ) assert p1.payouts is not None @@ -735,7 +734,7 @@ class TestProductFinancials: ledger_collection.initial_load(client=None, sync=True) pop_ledger_merge.build(client=client_no_amm, ledger_coll=ledger_collection) - p1.prebuild_balance( + p1.prebuild_balance_individual( thl_lm=thl_ledger_manager, ds=mnt_filepath, client=client_no_amm, @@ -750,7 +749,6 @@ class TestProductFinancials: assert p1.balance.available_balance == 70 p1.prebuild_payouts( - thl_lm=thl_ledger_manager, bp_pem=brokerage_product_payout_event_manager, ) assert p1.payouts is not None @@ -783,7 +781,7 @@ class TestProductFinancials: ledger_collection.initial_load(client=None, sync=True) pop_ledger_merge.build(client=client_no_amm, ledger_coll=ledger_collection) - p1.prebuild_balance( + p1.prebuild_balance_individual( thl_lm=thl_ledger_manager, ds=mnt_filepath, client=client_no_amm, @@ -798,7 +796,6 @@ class TestProductFinancials: assert p1.balance.available_balance == 66 p1.prebuild_payouts( - thl_lm=thl_ledger_manager, bp_pem=brokerage_product_payout_event_manager, ) assert p1.payouts is not None @@ -856,22 +853,17 @@ class TestProductFinancials: txs = thl_ledger_manager.get_tx_filtered_by_account(bp_wallet.uuid) assert len(txs) == 2 - with pytest.raises(expected_exception=AssertionError) as cm: - p1.prebuild_balance( - thl_lm=thl_ledger_manager, - ds=mnt_filepath, - client=client_no_amm, - ) - assert "Cannot build Product Balance" in str(cm.value) - ledger_collection.initial_load(client=None, sync=True) pop_ledger_merge.build(client=client_no_amm, ledger_coll=ledger_collection) - p1.prebuild_balance( + df = p1.get_pop_ledger_df( + product_ids=[p1.uuid], thl_lm=thl_ledger_manager, - ds=mnt_filepath, + pop_ledger=pop_ledger_merge, client=client_no_amm, ) + + p1.prebuild_balance(thl_lm=thl_ledger_manager, pop_ledger_df=df) assert isinstance(p1.balance, ProductBalances) assert p1.balance.payout == 114 assert p1.balance.adjustment == 0 @@ -881,19 +873,105 @@ class TestProductFinancials: body, content_type = p1.balance.to_prometheus() - p1.prebuild_private_balance( - thl_lm=thl_ledger_manager, - ds=mnt_filepath, + p1.prebuild_private_balance(thl_lm=thl_ledger_manager, pop_ledger_df=df) + assert p1.private_balance.commission == 5 * 2 + + p1.prebuild_user_wallet_balances(pop_ledger_df=df) + assert p1.user_wallet_balance.outstanding_liability == 38 * 2 + body, content_type = p1.user_wallet_balance.to_prometheus() + + def test_balance_both_kinds_of_products( + self, + gr_business: Business, + product_factory: Callable[..., Product], + user_factory: Callable[..., User], + mnt_filepath: GRLDatasets, + thl_ledger_manager: ThlLedgerManager, + start: datetime, + session_with_tx_factory: Callable[..., Session], + delete_ledger_db: Callable[..., None], + create_main_accounts: Callable[..., None], + client_no_amm: DaskClient, + ledger_collection: LedgerDFCollection, + pop_ledger_merge: PopLedgerMerge, + delete_df_collection: Callable[..., None], + payout_config, + thl_redis_config: RedisConfig, + brokerage_product_payout_event_manager: BrokerageProductPayoutEventManager, + ): + delete_ledger_db() + create_main_accounts() + delete_df_collection(coll=ledger_collection) + + p1: Product = product_factory( + business=gr_business, + ) + u1: User = user_factory(product=p1) + bp_wallet1 = thl_ledger_manager.get_account_or_create_bp_wallet(product=p1) + commission_wallet1 = thl_ledger_manager.get_account_or_create_bp_commission(p1) + p2: Product = product_factory( + business=gr_business, + user_wallet_config=UserWalletConfig(enabled=True), + payout_config=payout_config, + ) + u2: User = user_factory(product=p2) + bp_wallet2 = thl_ledger_manager.get_account_or_create_bp_wallet(product=p2) + user_wallet2 = thl_ledger_manager.get_account_or_create_user_wallet(user=u2) + commission_wallet2 = thl_ledger_manager.get_account_or_create_bp_commission(p2) + + session_with_tx_factory( + user=u1, + wall_req_cpi=Decimal("1.00"), + started=start, + ) + session_with_tx_factory( + user=u2, + wall_req_cpi=Decimal("1.00"), + started=start, + ) + + ledger_collection.initial_load(client=None, sync=True) + pop_ledger_merge.build(client=client_no_amm, ledger_coll=ledger_collection) + + df = Product.get_pop_ledger_df( + product_ids=[p1.uuid, p2.uuid], client=client_no_amm, + thl_lm=thl_ledger_manager, + pop_ledger=pop_ledger_merge, ) - assert p1.private_balance.commission == 5 * 2 + expected_accounts = [ + bp_wallet1.uuid, + bp_wallet2.uuid, + commission_wallet1.uuid, + commission_wallet2.uuid, + user_wallet2.uuid, + ] + assert set(df["account_id"]) == set(expected_accounts) + # The product_id column in pop ledger is set only for user txs + assert set(df["product_id"].dropna()) == {p2.uuid} - p1.prebuild_user_wallet_balances( - ds=mnt_filepath, + p1.prebuild_balance(thl_lm=thl_ledger_manager, pop_ledger_df=df) + assert p1.balance.payout == 95 + + p1.prebuild_private_balance(thl_lm=thl_ledger_manager, pop_ledger_df=df) + assert p1.private_balance.commission == 5 + + p2.prebuild_balance(thl_lm=thl_ledger_manager, pop_ledger_df=df) + assert p2.balance.payout == 57 + + p2.prebuild_private_balance(thl_lm=thl_ledger_manager, pop_ledger_df=df) + assert p2.private_balance.commission == 5 + + p2.prebuild_user_wallet_balances(pop_ledger_df=df) + assert p2.user_wallet_balance.outstanding_liability == 38 + + p1.set_cache( + thl_lm=thl_ledger_manager, client=client_no_amm, + bp_pem=brokerage_product_payout_event_manager, + redis_config=thl_redis_config, + pop_ledger=pop_ledger_merge, ) - assert p1.user_wallet_balance.outstanding_liability == 38 * 2 - body, content_type = p1.user_wallet_balance.to_prometheus() class TestProductBalance: @@ -1070,12 +1148,13 @@ class TestProductPOPFinancial: # --- test --- assert product.pop_financial is None - product.prebuild_pop_financial( + df = product.get_pop_ledger_df( + product_ids=[product.uuid], thl_lm=thl_ledger_manager, - ds=mnt_filepath, - client=client_no_amm, pop_ledger=pop_ledger_merge, + client=client_no_amm, ) + product.prebuild_pop_financial(thl_lm=thl_ledger_manager, pop_ledger_df=df) from generalresearch.models.thl.finance import POPFinancial @@ -1125,6 +1204,10 @@ class TestProductCache: create_main_accounts() delete_df_collection(coll=ledger_collection) + # In gr-api when we create a product, we create the bp wallet. + # maybe that should be standardized... + thl_ledger_manager.get_account_or_create_bp_wallet(product=product) + # Confirm the default / null behavior rc = thl_redis_config.create_redis_client() res: str | None = rc.get(product.cache_key) -- cgit v1.2.3