334 lines
12 KiB
Python
334 lines
12 KiB
Python
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))
|