add some helpers so we dont break existing tests. Add a test for the bulk load with both kinds of products.

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)