from datetime import datetime, timezone from io import StringIO from django.core.management import call_command from django.core.management.base import CommandError from django.test import TestCase from .brokerage_types import ( BrokerageAccountPayload, BrokerageAssetPayload, BrokerageHoldingPayload, BrokerageSnapshot, ) from .brokerages.base import BaseBrokerageAdapter, BrokerageAdapterError from .brokerages.random import RandomBrokerageAdapter from .brokerages.registry import register_brokerage_adapter from .models import Account, AccountHolding, Asset, BrokerageConnection, IngestionRun from .services import BrokerageIngestionService class FakeBrokerageAdapter(BaseBrokerageAdapter): provider = "fake-broker" def __init__(self, connection, snapshot=None, error=None): super().__init__(connection) self.snapshot = snapshot self.error = error def fetch_snapshot(self): if self.error: raise self.error return self.snapshot class AlternateFakeBrokerageAdapter(FakeBrokerageAdapter): provider = "alternate-fake-broker" class ExplodingBrokerageAdapter(BaseBrokerageAdapter): provider = "error-broker" def fetch_snapshot(self): raise RuntimeError("brokerage unavailable") class BrokerageIngestionServiceTests(TestCase): def setUp(self): register_brokerage_adapter(FakeBrokerageAdapter.provider, FakeBrokerageAdapter) register_brokerage_adapter( AlternateFakeBrokerageAdapter.provider, AlternateFakeBrokerageAdapter, ) register_brokerage_adapter( ExplodingBrokerageAdapter.provider, ExplodingBrokerageAdapter, ) self.connection = BrokerageConnection.objects.create( provider=FakeBrokerageAdapter.provider, name="Primary", base_url="https://broker.example/api/", credentials={"access_token": "token-1"}, settings={"account_endpoint": "/accounts"}, ) def test_ingest_creates_and_updates_portfolio_records(self): initial_snapshot = BrokerageSnapshot( accounts=[ BrokerageAccountPayload( external_id="acct-1", name="Brokerage Account", number="1234", balance=10000, ) ], assets=[ BrokerageAssetPayload( external_id="asset-1", ticker="AAPL", price=195.12, ) ], holdings=[ BrokerageHoldingPayload( external_id="holding-1", account_external_id="acct-1", asset_external_id="asset-1", ticker="AAPL", amount=10, cost_basis=150.25, purchase_date=datetime(2025, 1, 15, tzinfo=timezone.utc), ) ], ) updated_snapshot = BrokerageSnapshot( accounts=[ BrokerageAccountPayload( external_id="acct-1", name="Brokerage Account", number="1234", balance=12500, ) ], assets=[ BrokerageAssetPayload( external_id="asset-1", ticker="AAPL", price=205.55, ) ], holdings=[ BrokerageHoldingPayload( external_id="holding-1", account_external_id="acct-1", asset_external_id="asset-1", ticker="AAPL", amount=12, cost_basis=150.25, purchase_date=datetime(2025, 1, 15, tzinfo=timezone.utc), ) ], ) first_run = BrokerageIngestionService( self.connection, adapter=FakeBrokerageAdapter(self.connection, snapshot=initial_snapshot), ).ingest() second_run = BrokerageIngestionService( self.connection, adapter=FakeBrokerageAdapter(self.connection, snapshot=updated_snapshot), ).ingest() self.assertEqual(first_run.status, IngestionRun.STATUS_SUCCEEDED) self.assertEqual(second_run.status, IngestionRun.STATUS_SUCCEEDED) self.assertEqual(Account.objects.count(), 1) self.assertEqual(Asset.objects.count(), 1) self.assertEqual(AccountHolding.objects.count(), 1) account = Account.objects.get() asset = Asset.objects.get() holding = AccountHolding.objects.get() self.assertEqual(account.balance, 12500) self.assertEqual(asset.price, 205.55) self.assertEqual(holding.ammount, 12) self.assertEqual(second_run.accounts_processed, 1) self.assertEqual(second_run.assets_processed, 1) self.assertEqual(second_run.holdings_processed, 1) def test_ingest_marks_run_failed_when_client_errors(self): failing_connection = BrokerageConnection.objects.create( provider=ExplodingBrokerageAdapter.provider, name="Failing", base_url="https://broker.example/api/", ) with self.assertRaises(RuntimeError): BrokerageIngestionService(failing_connection).ingest() run = IngestionRun.objects.get() self.assertEqual(run.status, IngestionRun.STATUS_FAILED) self.assertIn("brokerage unavailable", run.error_message) def test_multiple_brokerages_can_ingest_same_external_ids(self): alternate_connection = BrokerageConnection.objects.create( provider=AlternateFakeBrokerageAdapter.provider, name="Secondary", base_url="https://broker-two.example/api/", ) snapshot = BrokerageSnapshot( accounts=[ BrokerageAccountPayload( external_id="acct-1", name="Shared External Account", number="9999", balance=5000, ) ], assets=[ BrokerageAssetPayload( external_id="asset-1", ticker="MSFT", price=410.01, ) ], holdings=[ BrokerageHoldingPayload( external_id="holding-1", account_external_id="acct-1", asset_external_id="asset-1", ticker="MSFT", amount=3, cost_basis=300.00, purchase_date=datetime(2025, 2, 1, tzinfo=timezone.utc), ) ], ) BrokerageIngestionService( self.connection, adapter=FakeBrokerageAdapter(self.connection, snapshot=snapshot), ).ingest() BrokerageIngestionService( alternate_connection, adapter=AlternateFakeBrokerageAdapter( alternate_connection, snapshot=snapshot, ), ).ingest() self.assertEqual(Account.objects.count(), 2) self.assertEqual(Asset.objects.count(), 2) self.assertEqual(AccountHolding.objects.count(), 2) self.assertEqual( set(Account.objects.values_list("provider", flat=True)), {FakeBrokerageAdapter.provider, AlternateFakeBrokerageAdapter.provider}, ) def test_missing_adapter_raises_clear_error(self): unknown_connection = BrokerageConnection.objects.create( provider="unknown-broker", name="Unknown", base_url="https://unknown.example/api/", ) with self.assertRaises(BrokerageAdapterError): BrokerageIngestionService(unknown_connection).ingest() def test_random_adapter_generates_seeded_snapshot(self): random_connection = BrokerageConnection.objects.create( provider=RandomBrokerageAdapter.provider, name="Randomized", base_url="https://local.random/", settings={ "seed": 17, "account_count": 2, "asset_count": 5, "holdings_per_account": 3, "tickers": ["AAPL", "MSFT", "TSLA", "NVDA", "AMD"], }, ) first_snapshot = RandomBrokerageAdapter(random_connection).fetch_snapshot() second_snapshot = RandomBrokerageAdapter(random_connection).fetch_snapshot() self.assertEqual(first_snapshot.accounts, second_snapshot.accounts) self.assertEqual(first_snapshot.assets, second_snapshot.assets) self.assertEqual(first_snapshot.holdings, second_snapshot.holdings) self.assertEqual(len(first_snapshot.accounts), 2) self.assertEqual(len(first_snapshot.assets), 5) self.assertEqual(len(first_snapshot.holdings), 6) def test_random_adapter_ingests_generated_data(self): random_connection = BrokerageConnection.objects.create( provider=RandomBrokerageAdapter.provider, name="Random Import", base_url="https://local.random/", settings={ "seed": 99, "account_count": 3, "asset_count": 6, "holdings_per_account": 2, }, ) run = BrokerageIngestionService(random_connection).ingest() self.assertEqual(run.status, IngestionRun.STATUS_SUCCEEDED) self.assertEqual(run.accounts_processed, 3) self.assertEqual(run.assets_processed, 6) self.assertEqual(run.holdings_processed, 6) self.assertEqual( Account.objects.filter(provider=RandomBrokerageAdapter.provider).count(), 3, ) self.assertEqual( Asset.objects.filter(provider=RandomBrokerageAdapter.provider).count(), 6, ) self.assertEqual( AccountHolding.objects.filter( provider=RandomBrokerageAdapter.provider ).count(), 6, ) def test_sync_command_accepts_provider_identifier(self): random_connection = BrokerageConnection.objects.create( provider=RandomBrokerageAdapter.provider, name="Demo Random", base_url="https://local.random/", settings={ "seed": 123, "account_count": 1, "asset_count": 2, "holdings_per_account": 1, }, ) output = StringIO() call_command( "sync_brokerage", connection=random_connection.provider, stdout=output, ) self.assertIn("Sync finished", output.getvalue()) self.assertEqual( Account.objects.filter(provider=RandomBrokerageAdapter.provider).count(), 1, ) def test_sync_command_requires_name_when_provider_has_multiple_connections(self): BrokerageConnection.objects.create( provider=RandomBrokerageAdapter.provider, name="Random One", base_url="https://local.random/", ) BrokerageConnection.objects.create( provider=RandomBrokerageAdapter.provider, name="Random Two", base_url="https://local.random/", ) with self.assertRaises(CommandError) as exc: call_command( "sync_brokerage", connection=RandomBrokerageAdapter.provider, ) self.assertIn("matches multiple active connections", str(exc.exception))