Stonks/overview/tests.py

617 lines
23 KiB
Python

from datetime import datetime, timedelta, 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 django.utils import timezone as django_timezone
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,
AccountSnapshot,
Asset,
AssetPriceSnapshot,
BrokerageConnection,
HoldingSnapshot,
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)
self.assertEqual(AccountSnapshot.objects.count(), 2)
self.assertEqual(AssetPriceSnapshot.objects.count(), 2)
self.assertEqual(HoldingSnapshot.objects.count(), 2)
account = Account.objects.get()
asset = Asset.objects.get()
holding = AccountHolding.objects.get()
latest_holding_snapshot = HoldingSnapshot.objects.filter(
ingestion_run=second_run
).get()
self.assertEqual(account.balance, 12500)
self.assertEqual(asset.price, 205.55)
self.assertEqual(holding.ammount, 12)
self.assertEqual(latest_holding_snapshot.value, 12 * 205.55)
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,
)
self.assertEqual(
AccountSnapshot.objects.filter(ingestion_run=run).count(),
3,
)
self.assertEqual(
AssetPriceSnapshot.objects.filter(ingestion_run=run).count(),
6,
)
self.assertEqual(
HoldingSnapshot.objects.filter(ingestion_run=run).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))
def test_overview_page_renders_portfolio_history(self):
BrokerageIngestionService(
self.connection,
adapter=FakeBrokerageAdapter(
self.connection,
snapshot=BrokerageSnapshot(
accounts=[
BrokerageAccountPayload(
external_id="acct-1",
name="Brokerage Account",
number="1234",
balance=10000,
)
],
assets=[
BrokerageAssetPayload(
external_id="asset-1",
ticker="AAPL",
price=200.00,
)
],
holdings=[
BrokerageHoldingPayload(
external_id="holding-1",
account_external_id="acct-1",
asset_external_id="asset-1",
ticker="AAPL",
amount=10,
cost_basis=150.00,
purchase_date=datetime(2025, 1, 15, tzinfo=timezone.utc),
)
],
),
),
).ingest()
response = self.client.get("/")
self.assertEqual(response.status_code, 200)
self.assertContains(response, "Portfolio Value Over Time")
self.assertContains(response, "Recent Sync Points")
def test_account_page_renders_account_history(self):
run = BrokerageIngestionService(
self.connection,
adapter=FakeBrokerageAdapter(
self.connection,
snapshot=BrokerageSnapshot(
accounts=[
BrokerageAccountPayload(
external_id="acct-1",
name="Brokerage Account",
number="1234",
balance=10000,
)
],
assets=[
BrokerageAssetPayload(
external_id="asset-1",
ticker="AAPL",
price=200.00,
)
],
holdings=[
BrokerageHoldingPayload(
external_id="holding-1",
account_external_id="acct-1",
asset_external_id="asset-1",
ticker="AAPL",
amount=10,
cost_basis=150.00,
purchase_date=datetime(2025, 1, 15, tzinfo=timezone.utc),
)
],
),
),
).ingest()
account = AccountSnapshot.objects.get(ingestion_run=run).account
response = self.client.get(f"/portfolio/{account.id}")
self.assertEqual(response.status_code, 200)
self.assertContains(response, "Account Value Over Time")
self.assertContains(response, "Current Holdings")
def test_overview_range_selector_filters_to_recent_history(self):
first_run = BrokerageIngestionService(
self.connection,
adapter=FakeBrokerageAdapter(
self.connection,
snapshot=BrokerageSnapshot(
accounts=[
BrokerageAccountPayload(
external_id="acct-1",
name="Brokerage Account",
number="1234",
balance=8000,
)
],
assets=[
BrokerageAssetPayload(
external_id="asset-1",
ticker="AAPL",
price=180.00,
)
],
holdings=[
BrokerageHoldingPayload(
external_id="holding-1",
account_external_id="acct-1",
asset_external_id="asset-1",
ticker="AAPL",
amount=10,
cost_basis=150.00,
purchase_date=datetime(2025, 1, 15, tzinfo=timezone.utc),
)
],
),
),
).ingest()
second_run = BrokerageIngestionService(
self.connection,
adapter=FakeBrokerageAdapter(
self.connection,
snapshot=BrokerageSnapshot(
accounts=[
BrokerageAccountPayload(
external_id="acct-1",
name="Brokerage Account",
number="1234",
balance=12000,
)
],
assets=[
BrokerageAssetPayload(
external_id="asset-1",
ticker="AAPL",
price=210.00,
)
],
holdings=[
BrokerageHoldingPayload(
external_id="holding-1",
account_external_id="acct-1",
asset_external_id="asset-1",
ticker="AAPL",
amount=10,
cost_basis=150.00,
purchase_date=datetime(2025, 1, 15, tzinfo=timezone.utc),
)
],
),
),
).ingest()
old_timestamp = django_timezone.now() - timedelta(days=40)
recent_timestamp = django_timezone.now() - timedelta(days=2)
IngestionRun.objects.filter(pk=first_run.pk).update(started_at=old_timestamp)
IngestionRun.objects.filter(pk=second_run.pk).update(started_at=recent_timestamp)
AccountSnapshot.objects.filter(ingestion_run=first_run).update(snapshot_at=old_timestamp)
AccountSnapshot.objects.filter(ingestion_run=second_run).update(snapshot_at=recent_timestamp)
response = self.client.get("/", {"range": "7d"})
self.assertEqual(response.status_code, 200)
self.assertContains(response, "value=\"7d\" selected")
self.assertNotContains(response, old_timestamp.strftime("%Y-%m-%d %H:%M"))
self.assertContains(response, recent_timestamp.strftime("%Y-%m-%d %H:%M"))
def test_account_range_selector_filters_account_history(self):
first_run = BrokerageIngestionService(
self.connection,
adapter=FakeBrokerageAdapter(
self.connection,
snapshot=BrokerageSnapshot(
accounts=[
BrokerageAccountPayload(
external_id="acct-1",
name="Brokerage Account",
number="1234",
balance=9000,
)
],
assets=[
BrokerageAssetPayload(
external_id="asset-1",
ticker="AAPL",
price=180.00,
)
],
holdings=[
BrokerageHoldingPayload(
external_id="holding-1",
account_external_id="acct-1",
asset_external_id="asset-1",
ticker="AAPL",
amount=10,
cost_basis=150.00,
purchase_date=datetime(2025, 1, 15, tzinfo=timezone.utc),
)
],
),
),
).ingest()
second_run = BrokerageIngestionService(
self.connection,
adapter=FakeBrokerageAdapter(
self.connection,
snapshot=BrokerageSnapshot(
accounts=[
BrokerageAccountPayload(
external_id="acct-1",
name="Brokerage Account",
number="1234",
balance=11000,
)
],
assets=[
BrokerageAssetPayload(
external_id="asset-1",
ticker="AAPL",
price=205.00,
)
],
holdings=[
BrokerageHoldingPayload(
external_id="holding-1",
account_external_id="acct-1",
asset_external_id="asset-1",
ticker="AAPL",
amount=10,
cost_basis=150.00,
purchase_date=datetime(2025, 1, 15, tzinfo=timezone.utc),
)
],
),
),
).ingest()
account = Account.objects.get()
old_timestamp = django_timezone.now() - timedelta(days=500)
recent_timestamp = django_timezone.now() - timedelta(days=5)
IngestionRun.objects.filter(pk=first_run.pk).update(started_at=old_timestamp)
IngestionRun.objects.filter(pk=second_run.pk).update(started_at=recent_timestamp)
AccountSnapshot.objects.filter(ingestion_run=first_run).update(snapshot_at=old_timestamp)
AccountSnapshot.objects.filter(ingestion_run=second_run).update(snapshot_at=recent_timestamp)
response = self.client.get(f"/portfolio/{account.id}", {"range": "30d"})
self.assertEqual(response.status_code, 200)
self.assertContains(response, "value=\"30d\" selected")
self.assertNotContains(response, old_timestamp.strftime("%Y-%m-%d %H:%M"))
self.assertContains(response, recent_timestamp.strftime("%Y-%m-%d %H:%M"))