AI: Brokerage Ingestion infra

This commit is contained in:
Tyler Wilcox 2026-05-23 21:18:05 -04:00
parent 2958a4b459
commit 618072069e
15 changed files with 1157 additions and 4 deletions

View File

@ -1,8 +1,10 @@
from django.contrib import admin from django.contrib import admin
from .models import Account, AccountHolding, Asset from .models import Account, AccountHolding, Asset, BrokerageConnection, IngestionRun
# Register your models here. # Register your models here.
admin.site.register(Account) admin.site.register(Account)
admin.site.register(AccountHolding) admin.site.register(AccountHolding)
admin.site.register(Asset) admin.site.register(Asset)
admin.site.register(BrokerageConnection)
admin.site.register(IngestionRun)

View File

@ -0,0 +1,35 @@
from dataclasses import dataclass, field
from datetime import datetime
@dataclass(frozen=True)
class BrokerageAccountPayload:
external_id: str
name: str
number: str
balance: int
@dataclass(frozen=True)
class BrokerageAssetPayload:
external_id: str
ticker: str
price: float
@dataclass(frozen=True)
class BrokerageHoldingPayload:
external_id: str
account_external_id: str
asset_external_id: str
ticker: str
amount: float
cost_basis: float
purchase_date: datetime
@dataclass(frozen=True)
class BrokerageSnapshot:
accounts: list[BrokerageAccountPayload] = field(default_factory=list)
assets: list[BrokerageAssetPayload] = field(default_factory=list)
holdings: list[BrokerageHoldingPayload] = field(default_factory=list)

View File

@ -0,0 +1,11 @@
from .base import BaseBrokerageAdapter, BrokerageAdapterError
from .registry import get_brokerage_adapter, register_brokerage_adapter
from .random import RandomBrokerageAdapter
__all__ = [
"BaseBrokerageAdapter",
"BrokerageAdapterError",
"RandomBrokerageAdapter",
"get_brokerage_adapter",
"register_brokerage_adapter",
]

View File

@ -0,0 +1,19 @@
from abc import ABC, abstractmethod
from overview.brokerage_types import BrokerageSnapshot
from overview.models import BrokerageConnection
class BrokerageAdapterError(Exception):
pass
class BaseBrokerageAdapter(ABC):
provider = ""
def __init__(self, connection: BrokerageConnection):
self.connection = connection
@abstractmethod
def fetch_snapshot(self) -> BrokerageSnapshot:
raise NotImplementedError

View File

@ -0,0 +1,172 @@
import random
from datetime import datetime, timedelta
from overview.brokerage_types import (
BrokerageAccountPayload,
BrokerageAssetPayload,
BrokerageHoldingPayload,
BrokerageSnapshot,
)
from .base import BaseBrokerageAdapter
class RandomBrokerageAdapter(BaseBrokerageAdapter):
provider = "random-demo"
DEFAULT_TICKERS = [
"AAPL",
"MSFT",
"GOOGL",
"AMZN",
"NVDA",
"TSLA",
"META",
"BRK.B",
"JPM",
"V",
"NFLX",
"AMD",
]
def fetch_snapshot(self) -> BrokerageSnapshot:
generator = random.Random(self.connection.get_setting("seed"))
account_count = max(1, int(self.connection.get_setting("account_count", 2)))
asset_count = max(1, int(self.connection.get_setting("asset_count", 8)))
holdings_per_account = max(
1,
int(self.connection.get_setting("holdings_per_account", 4)),
)
reference_timestamp = self.connection.get_setting(
"reference_timestamp",
"2025-01-01T00:00:00+00:00",
)
reference_time = datetime.fromisoformat(
str(reference_timestamp).replace("Z", "+00:00")
)
tickers = self.connection.get_setting("tickers", self.DEFAULT_TICKERS)
if not isinstance(tickers, list) or not tickers:
tickers = self.DEFAULT_TICKERS
assets = self._build_assets(generator, asset_count, tickers)
accounts = self._build_accounts(generator, account_count)
holdings = self._build_holdings(
generator,
accounts=accounts,
assets=assets,
holdings_per_account=holdings_per_account,
reference_time=reference_time,
)
return BrokerageSnapshot(
accounts=accounts,
assets=assets,
holdings=holdings,
)
def _build_accounts(
self,
generator: random.Random,
account_count: int,
) -> list[BrokerageAccountPayload]:
account_labels = [
"Individual",
"Joint",
"Retirement",
"Roth IRA",
"Trust",
"Education",
]
accounts = []
for index in range(account_count):
suffix = generator.randint(1000, 9999)
name = account_labels[index % len(account_labels)]
accounts.append(
BrokerageAccountPayload(
external_id=f"acct-{index + 1}",
name=f"{name} Account {index + 1}",
number=f"{suffix}{generator.randint(10, 99)}",
balance=0,
)
)
return accounts
def _build_assets(
self,
generator: random.Random,
asset_count: int,
tickers: list[str],
) -> list[BrokerageAssetPayload]:
assets = []
for index in range(asset_count):
ticker = tickers[index % len(tickers)]
base_price = round(generator.uniform(25, 750), 2)
assets.append(
BrokerageAssetPayload(
external_id=f"asset-{index + 1}",
ticker=ticker,
price=base_price,
)
)
return assets
def _build_holdings(
self,
generator: random.Random,
accounts: list[BrokerageAccountPayload],
assets: list[BrokerageAssetPayload],
holdings_per_account: int,
reference_time: datetime,
) -> list[BrokerageHoldingPayload]:
holdings = []
asset_lookup = {asset.external_id: asset for asset in assets}
days_back = int(self.connection.get_setting("max_purchase_age_days", 3650))
next_holding_id = 1
refreshed_accounts = []
for account in accounts:
selected_assets = generator.sample(
assets,
k=min(holdings_per_account, len(assets)),
)
balance = 0.0
for asset in selected_assets:
amount = round(generator.uniform(1, 50), 4)
current_value = amount * asset_lookup[asset.external_id].price
price_offset = generator.uniform(0.7, 1.15)
cost_basis = round(asset.price * price_offset, 2)
purchase_date = reference_time - timedelta(
days=generator.randint(1, max(1, days_back))
)
holdings.append(
BrokerageHoldingPayload(
external_id=f"holding-{next_holding_id}",
account_external_id=account.external_id,
asset_external_id=asset.external_id,
ticker=asset.ticker,
amount=amount,
cost_basis=cost_basis,
purchase_date=purchase_date,
)
)
next_holding_id += 1
balance += current_value
refreshed_accounts.append(
BrokerageAccountPayload(
external_id=account.external_id,
name=account.name,
number=account.number,
balance=int(round(balance)),
)
)
accounts[:] = refreshed_accounts
return holdings

View File

@ -0,0 +1,34 @@
from django.conf import settings
from django.utils.module_loading import import_string
from overview.models import BrokerageConnection
from .base import BaseBrokerageAdapter, BrokerageAdapterError
from .random import RandomBrokerageAdapter
from .rest import GenericRestBrokerageAdapter
_DEFAULT_ADAPTERS = {
GenericRestBrokerageAdapter.provider: GenericRestBrokerageAdapter,
RandomBrokerageAdapter.provider: RandomBrokerageAdapter,
}
def register_brokerage_adapter(provider: str, adapter_class: type[BaseBrokerageAdapter]):
_DEFAULT_ADAPTERS[provider] = adapter_class
def get_brokerage_adapter(connection: BrokerageConnection) -> BaseBrokerageAdapter:
adapter_class = _DEFAULT_ADAPTERS.get(connection.provider)
configured_paths = getattr(settings, "BROKERAGE_ADAPTERS", {})
configured_path = configured_paths.get(connection.provider)
if configured_path:
adapter_class = import_string(configured_path)
if not adapter_class:
raise BrokerageAdapterError(
f"No brokerage adapter is registered for provider '{connection.provider}'."
)
return adapter_class(connection)

131
overview/brokerages/rest.py Normal file
View File

@ -0,0 +1,131 @@
import json
from datetime import datetime, timezone
from urllib.parse import urljoin
from urllib.request import Request, urlopen
from overview.brokerage_types import (
BrokerageAccountPayload,
BrokerageAssetPayload,
BrokerageHoldingPayload,
BrokerageSnapshot,
)
from .base import BaseBrokerageAdapter, BrokerageAdapterError
class GenericRestBrokerageAdapter(BaseBrokerageAdapter):
provider = "generic-rest"
def fetch_snapshot(self) -> BrokerageSnapshot:
accounts = self._fetch_collection(
self.connection.get_setting("account_endpoint", "/accounts")
)
assets = self._fetch_collection(
self.connection.get_setting("asset_endpoint", "/assets")
)
holdings = self._fetch_collection(
self.connection.get_setting("holding_endpoint", "/holdings")
)
return BrokerageSnapshot(
accounts=[self._build_account(item) for item in accounts],
assets=[self._build_asset(item) for item in assets],
holdings=[self._build_holding(item) for item in holdings],
)
def _fetch_collection(self, endpoint: str) -> list[dict]:
url = urljoin(self.connection.base_url, endpoint)
request = Request(url, headers=self._headers())
try:
with urlopen(request, timeout=30) as response:
payload = json.load(response)
except Exception as exc:
raise BrokerageAdapterError(
f"Unable to fetch brokerage data from {url}: {exc}"
) from exc
if isinstance(payload, dict):
for key in ("results", "data", "items"):
value = payload.get(key)
if isinstance(value, list):
return value
if not isinstance(payload, list):
raise BrokerageAdapterError(
f"Expected a list payload from {url}, received {type(payload).__name__}"
)
return payload
def _headers(self) -> dict[str, str]:
headers = {"Accept": "application/json"}
token = self.connection.get_credential("access_token")
api_key = self.connection.get_credential("api_key")
api_key_header = self.connection.get_setting("api_key_header", "X-API-Key")
if token:
headers["Authorization"] = f"Bearer {token}"
if api_key:
headers[api_key_header] = api_key
extra_headers = self.connection.get_setting("headers", {})
if isinstance(extra_headers, dict):
headers.update({str(key): str(value) for key, value in extra_headers.items()})
return headers
def _build_account(self, item: dict) -> BrokerageAccountPayload:
account_id_field = self.connection.get_setting("account_id_field", "id")
account_name_field = self.connection.get_setting("account_name_field", "name")
account_number_field = self.connection.get_setting("account_number_field", "number")
account_balance_field = self.connection.get_setting("account_balance_field", "balance")
return BrokerageAccountPayload(
external_id=str(item[account_id_field]),
name=item[account_name_field],
number=item.get(account_number_field, ""),
balance=int(item.get(account_balance_field, 0)),
)
def _build_asset(self, item: dict) -> BrokerageAssetPayload:
asset_id_field = self.connection.get_setting("asset_id_field", "id")
asset_ticker_field = self.connection.get_setting("asset_ticker_field", "ticker")
asset_price_field = self.connection.get_setting("asset_price_field", "price")
return BrokerageAssetPayload(
external_id=str(item[asset_id_field]),
ticker=item[asset_ticker_field],
price=float(item.get(asset_price_field, 0)),
)
def _build_holding(self, item: dict) -> BrokerageHoldingPayload:
holding_id_field = self.connection.get_setting("holding_id_field", "id")
holding_account_field = self.connection.get_setting("holding_account_field", "account_id")
holding_asset_field = self.connection.get_setting("holding_asset_field", "asset_id")
holding_ticker_field = self.connection.get_setting("holding_ticker_field", "ticker")
holding_amount_field = self.connection.get_setting("holding_amount_field", "amount")
holding_cost_basis_field = self.connection.get_setting(
"holding_cost_basis_field", "cost_basis"
)
holding_purchase_date_field = self.connection.get_setting(
"holding_purchase_date_field", "purchase_date"
)
purchase_date = item.get(holding_purchase_date_field)
if purchase_date:
purchase_date = datetime.fromisoformat(
str(purchase_date).replace("Z", "+00:00")
)
else:
purchase_date = datetime.now(timezone.utc)
return BrokerageHoldingPayload(
external_id=str(item[holding_id_field]),
account_external_id=str(item[holding_account_field]),
asset_external_id=str(item[holding_asset_field]),
ticker=item[holding_ticker_field],
amount=float(item.get(holding_amount_field, 0)),
cost_basis=float(item.get(holding_cost_basis_field, 0)),
purchase_date=purchase_date,
)

View File

@ -0,0 +1 @@

View File

@ -0,0 +1 @@

View File

@ -0,0 +1,70 @@
from django.core.management.base import BaseCommand, CommandError
from overview.models import BrokerageConnection
from overview.services import BrokerageIngestionService
class Command(BaseCommand):
help = "Sync investment accounts, assets, and holdings from a brokerage REST API."
def add_arguments(self, parser):
parser.add_argument(
"--connection",
required=True,
help="Connection name or provider key to sync.",
)
def handle(self, *args, **options):
identifier = options["connection"]
connection = self._get_connection(identifier)
run = BrokerageIngestionService(connection).ingest()
self.stdout.write(
self.style.SUCCESS(
"Sync finished: "
f"{run.accounts_processed} account(s), "
f"{run.assets_processed} asset(s), "
f"{run.holdings_processed} holding(s)."
)
)
def _get_connection(self, identifier: str) -> BrokerageConnection:
named_connection = BrokerageConnection.objects.filter(
name=identifier,
is_active=True,
).first()
if named_connection:
return named_connection
provider_matches = list(
BrokerageConnection.objects.filter(
provider=identifier,
is_active=True,
).order_by("name")
)
if len(provider_matches) == 1:
return provider_matches[0]
if len(provider_matches) > 1:
matches = ", ".join(connection.name for connection in provider_matches)
raise CommandError(
f"Provider '{identifier}' matches multiple active connections: {matches}. "
"Re-run with --connection set to the connection name."
)
active_connections = list(
BrokerageConnection.objects.filter(is_active=True).order_by(
"provider",
"name",
)
)
available = ", ".join(
f"{connection.name} ({connection.provider})"
for connection in active_connections
)
if not available:
available = "none"
raise CommandError(
f"Active brokerage connection or provider '{identifier}' was not found. "
f"Available active connections: {available}."
)

View File

@ -0,0 +1,96 @@
# Generated by Django 6.0.5 on 2026-05-24 00:51
import django.db.models.deletion
from django.db import migrations, models
class Migration(migrations.Migration):
dependencies = [
('overview', '0003_accountholding_delete_assetholding_asset_price_and_more'),
]
operations = [
migrations.CreateModel(
name='BrokerageConnection',
fields=[
('id', models.BigAutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')),
('provider', models.CharField(max_length=64)),
('name', models.CharField(max_length=128)),
('base_url', models.URLField()),
('access_token', models.CharField(blank=True, max_length=512)),
('api_key', models.CharField(blank=True, max_length=255)),
('account_endpoint', models.CharField(default='/accounts', max_length=255)),
('asset_endpoint', models.CharField(default='/assets', max_length=255)),
('holding_endpoint', models.CharField(default='/holdings', max_length=255)),
('is_active', models.BooleanField(default=True)),
('created_at', models.DateTimeField(auto_now_add=True)),
('updated_at', models.DateTimeField(auto_now=True)),
],
),
migrations.CreateModel(
name='IngestionRun',
fields=[
('id', models.BigAutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')),
('status', models.CharField(choices=[('pending', 'Pending'), ('running', 'Running'), ('succeeded', 'Succeeded'), ('failed', 'Failed')], default='pending', max_length=16)),
('started_at', models.DateTimeField(blank=True, null=True)),
('finished_at', models.DateTimeField(blank=True, null=True)),
('accounts_processed', models.PositiveIntegerField(default=0)),
('assets_processed', models.PositiveIntegerField(default=0)),
('holdings_processed', models.PositiveIntegerField(default=0)),
('error_message', models.TextField(blank=True)),
('created_at', models.DateTimeField(auto_now_add=True)),
],
),
migrations.AddField(
model_name='account',
name='external_id',
field=models.CharField(blank=True, max_length=255),
),
migrations.AddField(
model_name='account',
name='provider',
field=models.CharField(blank=True, max_length=64),
),
migrations.AddField(
model_name='accountholding',
name='external_id',
field=models.CharField(blank=True, max_length=255),
),
migrations.AddField(
model_name='accountholding',
name='provider',
field=models.CharField(blank=True, max_length=64),
),
migrations.AddField(
model_name='asset',
name='external_id',
field=models.CharField(blank=True, max_length=255),
),
migrations.AddField(
model_name='asset',
name='provider',
field=models.CharField(blank=True, max_length=64),
),
migrations.AddConstraint(
model_name='account',
constraint=models.UniqueConstraint(fields=('provider', 'external_id'), name='unique_account_provider_external_id'),
),
migrations.AddConstraint(
model_name='accountholding',
constraint=models.UniqueConstraint(fields=('provider', 'external_id'), name='unique_holding_provider_external_id'),
),
migrations.AddConstraint(
model_name='asset',
constraint=models.UniqueConstraint(fields=('provider', 'external_id'), name='unique_asset_provider_external_id'),
),
migrations.AddConstraint(
model_name='brokerageconnection',
constraint=models.UniqueConstraint(fields=('provider', 'name'), name='unique_brokerage_connection_provider_name'),
),
migrations.AddField(
model_name='ingestionrun',
name='connection',
field=models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, to='overview.brokerageconnection'),
),
]

View File

@ -0,0 +1,43 @@
# Generated by Django 6.0.5 on 2026-05-24 00:58
from django.db import migrations, models
class Migration(migrations.Migration):
dependencies = [
('overview', '0004_brokerageconnection_ingestionrun_account_external_id_and_more'),
]
operations = [
migrations.RemoveField(
model_name='brokerageconnection',
name='access_token',
),
migrations.RemoveField(
model_name='brokerageconnection',
name='account_endpoint',
),
migrations.RemoveField(
model_name='brokerageconnection',
name='api_key',
),
migrations.RemoveField(
model_name='brokerageconnection',
name='asset_endpoint',
),
migrations.RemoveField(
model_name='brokerageconnection',
name='holding_endpoint',
),
migrations.AddField(
model_name='brokerageconnection',
name='credentials',
field=models.JSONField(blank=True, default=dict),
),
migrations.AddField(
model_name='brokerageconnection',
name='settings',
field=models.JSONField(blank=True, default=dict),
),
]

View File

@ -1,14 +1,40 @@
from django.db import models from django.db import models
# Create your models here.
class Account(models.Model): class Account(models.Model):
name = models.TextField(max_length=1024) name = models.TextField(max_length=1024)
number = models.TextField(max_length=1024) number = models.TextField(max_length=1024)
balance = models.IntegerField() balance = models.IntegerField()
provider = models.CharField(max_length=64, blank=True)
external_id = models.CharField(max_length=255, blank=True)
class Meta:
constraints = [
models.UniqueConstraint(
fields=["provider", "external_id"],
name="unique_account_provider_external_id",
),
]
def __str__(self):
return self.name
class Asset(models.Model): class Asset(models.Model):
ticker = models.TextField(max_length=20) ticker = models.TextField(max_length=20)
price = models.FloatField() price = models.FloatField()
provider = models.CharField(max_length=64, blank=True)
external_id = models.CharField(max_length=255, blank=True)
class Meta:
constraints = [
models.UniqueConstraint(
fields=["provider", "external_id"],
name="unique_asset_provider_external_id",
),
]
def __str__(self):
return self.ticker
class AccountHolding(models.Model): class AccountHolding(models.Model):
account = models.ForeignKey(Account, on_delete=models.CASCADE) account = models.ForeignKey(Account, on_delete=models.CASCADE)
@ -17,6 +43,16 @@ class AccountHolding(models.Model):
ammount = models.FloatField() ammount = models.FloatField()
cost_basis = models.FloatField() cost_basis = models.FloatField()
purchase_date = models.DateTimeField() purchase_date = models.DateTimeField()
provider = models.CharField(max_length=64, blank=True)
external_id = models.CharField(max_length=255, blank=True)
class Meta:
constraints = [
models.UniqueConstraint(
fields=["provider", "external_id"],
name="unique_holding_provider_external_id",
),
]
def holding_value(self): def holding_value(self):
return self.ammount * self.asset.price return self.ammount * self.asset.price
@ -27,3 +63,64 @@ class AccountHolding(models.Model):
def holding_gain(self): def holding_gain(self):
return self.holding_value() - self.holding_principle() return self.holding_value() - self.holding_principle()
def __str__(self):
return f"{self.account.name}: {self.ticker}"
class BrokerageConnection(models.Model):
provider = models.CharField(max_length=64)
name = models.CharField(max_length=128)
base_url = models.URLField()
credentials = models.JSONField(default=dict, blank=True)
settings = models.JSONField(default=dict, blank=True)
is_active = models.BooleanField(default=True)
created_at = models.DateTimeField(auto_now_add=True)
updated_at = models.DateTimeField(auto_now=True)
class Meta:
constraints = [
models.UniqueConstraint(
fields=["provider", "name"],
name="unique_brokerage_connection_provider_name",
),
]
def __str__(self):
return f"{self.provider}: {self.name}"
def get_credential(self, key: str, default=None):
return self.credentials.get(key, default)
def get_setting(self, key: str, default=None):
return self.settings.get(key, default)
class IngestionRun(models.Model):
STATUS_PENDING = "pending"
STATUS_RUNNING = "running"
STATUS_SUCCEEDED = "succeeded"
STATUS_FAILED = "failed"
STATUS_CHOICES = [
(STATUS_PENDING, "Pending"),
(STATUS_RUNNING, "Running"),
(STATUS_SUCCEEDED, "Succeeded"),
(STATUS_FAILED, "Failed"),
]
connection = models.ForeignKey(BrokerageConnection, on_delete=models.CASCADE)
status = models.CharField(
max_length=16,
choices=STATUS_CHOICES,
default=STATUS_PENDING,
)
started_at = models.DateTimeField(null=True, blank=True)
finished_at = models.DateTimeField(null=True, blank=True)
accounts_processed = models.PositiveIntegerField(default=0)
assets_processed = models.PositiveIntegerField(default=0)
holdings_processed = models.PositiveIntegerField(default=0)
error_message = models.TextField(blank=True)
created_at = models.DateTimeField(auto_now_add=True)
def __str__(self):
return f"{self.connection} [{self.status}]"

111
overview/services.py Normal file
View File

@ -0,0 +1,111 @@
from dataclasses import dataclass
from django.db import transaction
from django.utils import timezone
from .brokerage_types import BrokerageSnapshot
from .brokerages.base import BaseBrokerageAdapter
from .brokerages.registry import get_brokerage_adapter
from .models import Account, AccountHolding, Asset, BrokerageConnection, IngestionRun
@dataclass(frozen=True)
class IngestionResult:
accounts_processed: int
assets_processed: int
holdings_processed: int
class BrokerageIngestionService:
def __init__(
self,
connection: BrokerageConnection,
adapter: BaseBrokerageAdapter | None = None,
):
self.connection = connection
self.adapter = adapter or get_brokerage_adapter(connection)
def ingest(self) -> IngestionRun:
run = IngestionRun.objects.create(
connection=self.connection,
status=IngestionRun.STATUS_RUNNING,
started_at=timezone.now(),
)
try:
snapshot = self.adapter.fetch_snapshot()
result = self._persist_snapshot(snapshot)
except Exception as exc:
run.status = IngestionRun.STATUS_FAILED
run.error_message = str(exc)
run.finished_at = timezone.now()
run.save(
update_fields=["status", "error_message", "finished_at"]
)
raise
run.status = IngestionRun.STATUS_SUCCEEDED
run.finished_at = timezone.now()
run.accounts_processed = result.accounts_processed
run.assets_processed = result.assets_processed
run.holdings_processed = result.holdings_processed
run.save(
update_fields=[
"status",
"finished_at",
"accounts_processed",
"assets_processed",
"holdings_processed",
]
)
return run
@transaction.atomic
def _persist_snapshot(self, snapshot: BrokerageSnapshot) -> IngestionResult:
accounts_by_external_id = {}
assets_by_external_id = {}
for payload in snapshot.accounts:
account, _ = Account.objects.update_or_create(
provider=self.connection.provider,
external_id=payload.external_id,
defaults={
"name": payload.name,
"number": payload.number,
"balance": payload.balance,
},
)
accounts_by_external_id[payload.external_id] = account
for payload in snapshot.assets:
asset, _ = Asset.objects.update_or_create(
provider=self.connection.provider,
external_id=payload.external_id,
defaults={
"ticker": payload.ticker,
"price": payload.price,
},
)
assets_by_external_id[payload.external_id] = asset
for payload in snapshot.holdings:
account = accounts_by_external_id[payload.account_external_id]
asset = assets_by_external_id[payload.asset_external_id]
AccountHolding.objects.update_or_create(
provider=self.connection.provider,
external_id=payload.external_id,
defaults={
"account": account,
"asset": asset,
"ticker": payload.ticker,
"ammount": payload.amount,
"cost_basis": payload.cost_basis,
"purchase_date": payload.purchase_date,
},
)
return IngestionResult(
accounts_processed=len(snapshot.accounts),
assets_processed=len(snapshot.assets),
holdings_processed=len(snapshot.holdings),
)

View File

@ -1,3 +1,333 @@
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 django.test import TestCase
# Create your tests here. 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))