parent
9f00a1d0d2
commit
cdcdabecab
@ -0,0 +1,227 @@
|
||||
"""
|
||||
Balance tracking model for wallet balance snapshots at significant events.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from collections.abc import Sequence
|
||||
from datetime import datetime
|
||||
from enum import Enum
|
||||
from typing import Any, ClassVar
|
||||
|
||||
from sqlalchemy import Float, ForeignKey, Integer, String, desc, select
|
||||
from sqlalchemy.orm import Mapped, mapped_column
|
||||
|
||||
from freqtrade.persistence.base import ModelBase, SessionType
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class BalanceEventType(str, Enum):
|
||||
"""Types of events that trigger balance snapshots."""
|
||||
|
||||
BOT_START = "bot_start"
|
||||
TRADE_OPEN = "trade_open"
|
||||
TRADE_CLOSE = "trade_close"
|
||||
ORDER_FILL = "order_fill"
|
||||
ORDER_CANCEL = "order_cancel"
|
||||
POSITION_ADJUST = "position_adjust" # DCA
|
||||
SCHEDULED = "scheduled"
|
||||
MANUAL = "manual"
|
||||
|
||||
|
||||
class Balance(ModelBase):
|
||||
"""
|
||||
Balance database model.
|
||||
Tracks wallet balance snapshots at significant events to enable:
|
||||
- Detection of deposits/withdrawals
|
||||
- Accurate long-term PnL tracking
|
||||
- Historical balance reconstruction
|
||||
"""
|
||||
|
||||
__tablename__ = "balances"
|
||||
__allow_unmapped__ = True
|
||||
session: ClassVar[SessionType]
|
||||
|
||||
id: Mapped[int] = mapped_column(Integer, primary_key=True)
|
||||
timestamp: Mapped[datetime] = mapped_column(nullable=False, index=True)
|
||||
event_type: Mapped[str] = mapped_column(String(50), nullable=False, index=True)
|
||||
|
||||
# Balance data (stake currency)
|
||||
total_balance: Mapped[float] = mapped_column(Float(), nullable=False)
|
||||
free_balance: Mapped[float] = mapped_column(Float(), nullable=False)
|
||||
used_balance: Mapped[float | None] = mapped_column(Float(), nullable=True)
|
||||
stake_currency: Mapped[str] = mapped_column(String(25), nullable=False)
|
||||
|
||||
# Trade/Order association (optional)
|
||||
ft_trade_id: Mapped[int | None] = mapped_column(
|
||||
Integer, ForeignKey("trades.id"), nullable=True, index=True
|
||||
)
|
||||
ft_order_id: Mapped[int | None] = mapped_column(Integer, ForeignKey("orders.id"), nullable=True)
|
||||
|
||||
# Profit metrics at this point
|
||||
total_profit: Mapped[float | None] = mapped_column(Float(), nullable=True)
|
||||
open_trade_value: Mapped[float | None] = mapped_column(Float(), nullable=True)
|
||||
|
||||
# Deposit/withdrawal detection
|
||||
external_change: Mapped[float | None] = mapped_column(Float(), nullable=True)
|
||||
|
||||
is_dry_run: Mapped[bool] = mapped_column(nullable=False, default=False)
|
||||
context: Mapped[str | None] = mapped_column(String(1024), nullable=True)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return (
|
||||
f"Balance(id={self.id}, timestamp={self.timestamp}, "
|
||||
f"event_type={self.event_type}, total={self.total_balance}, "
|
||||
f"external_change={self.external_change})"
|
||||
)
|
||||
|
||||
def to_json(self) -> dict[str, Any]:
|
||||
"""Convert balance record to JSON-serializable dict."""
|
||||
return {
|
||||
"id": self.id,
|
||||
"timestamp": self.timestamp.isoformat() if self.timestamp else None,
|
||||
"event_type": self.event_type,
|
||||
"total_balance": self.total_balance,
|
||||
"free_balance": self.free_balance,
|
||||
"used_balance": self.used_balance,
|
||||
"stake_currency": self.stake_currency,
|
||||
"ft_trade_id": self.ft_trade_id,
|
||||
"ft_order_id": self.ft_order_id,
|
||||
"total_profit": self.total_profit,
|
||||
"open_trade_value": self.open_trade_value,
|
||||
"external_change": self.external_change,
|
||||
"is_dry_run": self.is_dry_run,
|
||||
"context": self.context,
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def get_balances(
|
||||
*,
|
||||
event_type: BalanceEventType | str | None = None,
|
||||
trade_id: int | None = None,
|
||||
since: datetime | None = None,
|
||||
until: datetime | None = None,
|
||||
limit: int | None = None,
|
||||
) -> Sequence["Balance"]:
|
||||
"""
|
||||
Query balance records with optional filters.
|
||||
|
||||
:param event_type: Filter by event type
|
||||
:param trade_id: Filter by trade ID
|
||||
:param since: Filter records after this timestamp
|
||||
:param until: Filter records before this timestamp
|
||||
:param limit: Maximum number of records to return
|
||||
:return: List of Balance records
|
||||
"""
|
||||
filters = []
|
||||
if event_type:
|
||||
event_str = event_type.value if isinstance(event_type, BalanceEventType) else event_type
|
||||
filters.append(Balance.event_type == event_str)
|
||||
if trade_id is not None:
|
||||
filters.append(Balance.ft_trade_id == trade_id)
|
||||
if since:
|
||||
filters.append(Balance.timestamp >= since)
|
||||
if until:
|
||||
filters.append(Balance.timestamp <= until)
|
||||
|
||||
query = select(Balance).filter(*filters).order_by(desc(Balance.timestamp))
|
||||
|
||||
if limit:
|
||||
query = query.limit(limit)
|
||||
|
||||
return Balance.session.scalars(query).all()
|
||||
|
||||
@staticmethod
|
||||
def get_latest(stake_currency: str | None = None) -> "Balance | None":
|
||||
"""
|
||||
Get the most recent balance record.
|
||||
|
||||
:param stake_currency: Optionally filter by stake currency
|
||||
:return: Latest Balance record or None
|
||||
"""
|
||||
filters = []
|
||||
if stake_currency:
|
||||
filters.append(Balance.stake_currency == stake_currency)
|
||||
|
||||
query = select(Balance).filter(*filters).order_by(desc(Balance.timestamp)).limit(1)
|
||||
return Balance.session.scalars(query).first()
|
||||
|
||||
@staticmethod
|
||||
def get_latest_before(
|
||||
timestamp: datetime, stake_currency: str | None = None
|
||||
) -> "Balance | None":
|
||||
"""
|
||||
Get the most recent balance record before the given timestamp.
|
||||
|
||||
:param timestamp: Get balance before this time
|
||||
:param stake_currency: Optionally filter by stake currency
|
||||
:return: Balance record or None
|
||||
"""
|
||||
filters = [Balance.timestamp < timestamp]
|
||||
if stake_currency:
|
||||
filters.append(Balance.stake_currency == stake_currency)
|
||||
|
||||
query = select(Balance).filter(*filters).order_by(desc(Balance.timestamp)).limit(1)
|
||||
return Balance.session.scalars(query).first()
|
||||
|
||||
@staticmethod
|
||||
def calculate_external_change(current: "Balance") -> float | None:
|
||||
"""
|
||||
Calculate external change (deposit/withdrawal) between previous and current balance.
|
||||
|
||||
External change = actual balance - expected balance
|
||||
Expected = previous_total + profit_change
|
||||
|
||||
:param current: Current balance record
|
||||
:return: External change amount or None if no previous record exists
|
||||
"""
|
||||
previous = Balance.get_latest_before(current.timestamp, current.stake_currency)
|
||||
if not previous:
|
||||
return None
|
||||
|
||||
# Expected = previous_total + profit_change
|
||||
profit_change = (current.total_profit or 0) - (previous.total_profit or 0)
|
||||
expected_balance = previous.total_balance + profit_change
|
||||
|
||||
# External change = actual - expected
|
||||
external_change = current.total_balance - expected_balance
|
||||
|
||||
# Ignore floating point noise (< 0.01% of balance)
|
||||
if current.total_balance > 0 and abs(external_change) > current.total_balance * 0.0001:
|
||||
return external_change
|
||||
return 0.0
|
||||
|
||||
@staticmethod
|
||||
def detect_deposits_withdrawals(
|
||||
since: datetime | None = None,
|
||||
until: datetime | None = None,
|
||||
threshold: float = 0.0,
|
||||
) -> list[dict[str, Any]]:
|
||||
"""
|
||||
Detect deposits and withdrawals by finding balance records with significant
|
||||
external changes.
|
||||
|
||||
:param since: Start of time range
|
||||
:param until: End of time range
|
||||
:param threshold: Minimum absolute change to report (default 0.0)
|
||||
:return: List of records with external changes
|
||||
"""
|
||||
balances = Balance.get_balances(since=since, until=until)
|
||||
|
||||
changes = []
|
||||
for balance in balances:
|
||||
if balance.external_change is not None and abs(balance.external_change) > threshold:
|
||||
change_type = "deposit" if balance.external_change > 0 else "withdrawal"
|
||||
changes.append(
|
||||
{
|
||||
"timestamp": balance.timestamp.isoformat() if balance.timestamp else None,
|
||||
"type": change_type,
|
||||
"amount": balance.external_change,
|
||||
"event_type": balance.event_type,
|
||||
"total_balance": balance.total_balance,
|
||||
"balance_id": balance.id,
|
||||
}
|
||||
)
|
||||
|
||||
return changes
|
||||
@ -0,0 +1,441 @@
|
||||
# pragma pylint: disable=missing-docstring, C0103
|
||||
"""
|
||||
Tests for the Balance model and BalanceEventType enum.
|
||||
"""
|
||||
|
||||
from datetime import UTC, datetime, timedelta
|
||||
|
||||
import pytest
|
||||
|
||||
from freqtrade.persistence import Balance, BalanceEventType, Trade
|
||||
from freqtrade.util import dt_now
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("init_persistence")
|
||||
class TestBalanceEventType:
|
||||
def test_event_type_values(self):
|
||||
"""Test that all expected event types exist"""
|
||||
assert BalanceEventType.BOT_START.value == "bot_start"
|
||||
assert BalanceEventType.TRADE_OPEN.value == "trade_open"
|
||||
assert BalanceEventType.TRADE_CLOSE.value == "trade_close"
|
||||
assert BalanceEventType.ORDER_FILL.value == "order_fill"
|
||||
assert BalanceEventType.ORDER_CANCEL.value == "order_cancel"
|
||||
assert BalanceEventType.POSITION_ADJUST.value == "position_adjust"
|
||||
assert BalanceEventType.SCHEDULED.value == "scheduled"
|
||||
assert BalanceEventType.MANUAL.value == "manual"
|
||||
|
||||
def test_event_type_is_str(self):
|
||||
"""Test that BalanceEventType inherits from str"""
|
||||
assert isinstance(BalanceEventType.BOT_START, str)
|
||||
assert BalanceEventType.BOT_START == "bot_start"
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("init_persistence")
|
||||
class TestBalanceModel:
|
||||
def test_balance_creation(self):
|
||||
"""Test creating a balance record"""
|
||||
balance = Balance(
|
||||
timestamp=dt_now(),
|
||||
event_type=BalanceEventType.BOT_START.value,
|
||||
total_balance=1000.0,
|
||||
free_balance=900.0,
|
||||
used_balance=100.0,
|
||||
stake_currency="USDT",
|
||||
is_dry_run=True,
|
||||
)
|
||||
Balance.session.add(balance)
|
||||
Trade.commit()
|
||||
|
||||
assert balance.id is not None
|
||||
assert balance.total_balance == 1000.0
|
||||
assert balance.free_balance == 900.0
|
||||
assert balance.used_balance == 100.0
|
||||
assert balance.stake_currency == "USDT"
|
||||
assert balance.is_dry_run is True
|
||||
|
||||
def test_balance_to_json(self):
|
||||
"""Test converting balance to JSON"""
|
||||
now = dt_now()
|
||||
balance = Balance(
|
||||
timestamp=now,
|
||||
event_type=BalanceEventType.TRADE_CLOSE.value,
|
||||
total_balance=1500.0,
|
||||
free_balance=1400.0,
|
||||
used_balance=100.0,
|
||||
stake_currency="USDT",
|
||||
total_profit=500.0,
|
||||
open_trade_value=100.0,
|
||||
external_change=0.0,
|
||||
is_dry_run=False,
|
||||
)
|
||||
Balance.session.add(balance)
|
||||
Trade.commit()
|
||||
|
||||
json_data = balance.to_json()
|
||||
assert json_data["id"] == balance.id
|
||||
assert json_data["event_type"] == "trade_close"
|
||||
assert json_data["total_balance"] == 1500.0
|
||||
assert json_data["free_balance"] == 1400.0
|
||||
assert json_data["used_balance"] == 100.0
|
||||
assert json_data["stake_currency"] == "USDT"
|
||||
assert json_data["total_profit"] == 500.0
|
||||
assert json_data["open_trade_value"] == 100.0
|
||||
assert json_data["external_change"] == 0.0
|
||||
assert json_data["is_dry_run"] is False
|
||||
|
||||
def test_balance_repr(self):
|
||||
"""Test balance string representation"""
|
||||
balance = Balance(
|
||||
timestamp=dt_now(),
|
||||
event_type=BalanceEventType.BOT_START.value,
|
||||
total_balance=1000.0,
|
||||
free_balance=900.0,
|
||||
stake_currency="USDT",
|
||||
is_dry_run=True,
|
||||
)
|
||||
Balance.session.add(balance)
|
||||
Trade.commit()
|
||||
|
||||
repr_str = repr(balance)
|
||||
assert "Balance" in repr_str
|
||||
assert "bot_start" in repr_str
|
||||
assert "total=1000.0" in repr_str
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("init_persistence")
|
||||
class TestBalanceQueries:
|
||||
def _create_test_balances(self):
|
||||
"""Create test balance records"""
|
||||
base_time = datetime(2024, 1, 1, 12, 0, 0, tzinfo=UTC)
|
||||
balances = [
|
||||
Balance(
|
||||
timestamp=base_time,
|
||||
event_type=BalanceEventType.BOT_START.value,
|
||||
total_balance=1000.0,
|
||||
free_balance=1000.0,
|
||||
stake_currency="USDT",
|
||||
total_profit=0.0,
|
||||
is_dry_run=True,
|
||||
),
|
||||
Balance(
|
||||
timestamp=base_time + timedelta(hours=1),
|
||||
event_type=BalanceEventType.TRADE_OPEN.value,
|
||||
total_balance=900.0,
|
||||
free_balance=800.0,
|
||||
used_balance=100.0,
|
||||
stake_currency="USDT",
|
||||
total_profit=0.0,
|
||||
ft_trade_id=1,
|
||||
is_dry_run=True,
|
||||
),
|
||||
Balance(
|
||||
timestamp=base_time + timedelta(hours=2),
|
||||
event_type=BalanceEventType.TRADE_CLOSE.value,
|
||||
total_balance=1050.0,
|
||||
free_balance=1050.0,
|
||||
stake_currency="USDT",
|
||||
total_profit=50.0,
|
||||
ft_trade_id=1,
|
||||
is_dry_run=True,
|
||||
),
|
||||
Balance(
|
||||
timestamp=base_time + timedelta(hours=3),
|
||||
event_type=BalanceEventType.BOT_START.value,
|
||||
total_balance=1200.0, # External deposit of ~150
|
||||
free_balance=1200.0,
|
||||
stake_currency="USDT",
|
||||
total_profit=50.0,
|
||||
external_change=150.0,
|
||||
is_dry_run=True,
|
||||
),
|
||||
]
|
||||
for b in balances:
|
||||
Balance.session.add(b)
|
||||
Trade.commit()
|
||||
return balances
|
||||
|
||||
def test_get_balances_all(self):
|
||||
"""Test getting all balance records"""
|
||||
self._create_test_balances()
|
||||
balances = Balance.get_balances()
|
||||
assert len(balances) == 4
|
||||
|
||||
def test_get_balances_by_event_type(self):
|
||||
"""Test filtering by event type"""
|
||||
self._create_test_balances()
|
||||
balances = Balance.get_balances(event_type=BalanceEventType.BOT_START)
|
||||
assert len(balances) == 2
|
||||
assert all(b.event_type == "bot_start" for b in balances)
|
||||
|
||||
def test_get_balances_by_event_type_string(self):
|
||||
"""Test filtering by event type as string"""
|
||||
self._create_test_balances()
|
||||
balances = Balance.get_balances(event_type="trade_close")
|
||||
assert len(balances) == 1
|
||||
assert balances[0].event_type == "trade_close"
|
||||
|
||||
def test_get_balances_by_trade_id(self):
|
||||
"""Test filtering by trade ID"""
|
||||
self._create_test_balances()
|
||||
balances = Balance.get_balances(trade_id=1)
|
||||
assert len(balances) == 2
|
||||
assert all(b.ft_trade_id == 1 for b in balances)
|
||||
|
||||
def test_get_balances_by_time_range(self):
|
||||
"""Test filtering by time range"""
|
||||
self._create_test_balances()
|
||||
base_time = datetime(2024, 1, 1, 12, 0, 0, tzinfo=UTC)
|
||||
balances = Balance.get_balances(
|
||||
since=base_time + timedelta(minutes=30),
|
||||
until=base_time + timedelta(hours=2, minutes=30),
|
||||
)
|
||||
assert len(balances) == 2
|
||||
assert all(b.event_type in ["trade_open", "trade_close"] for b in balances)
|
||||
|
||||
def test_get_balances_with_limit(self):
|
||||
"""Test limiting results"""
|
||||
self._create_test_balances()
|
||||
balances = Balance.get_balances(limit=2)
|
||||
assert len(balances) == 2
|
||||
|
||||
def test_get_latest(self):
|
||||
"""Test getting the most recent balance"""
|
||||
self._create_test_balances()
|
||||
latest = Balance.get_latest()
|
||||
assert latest is not None
|
||||
assert latest.total_balance == 1200.0
|
||||
|
||||
def test_get_latest_by_currency(self):
|
||||
"""Test getting latest balance by currency"""
|
||||
self._create_test_balances()
|
||||
latest = Balance.get_latest(stake_currency="USDT")
|
||||
assert latest is not None
|
||||
assert latest.stake_currency == "USDT"
|
||||
|
||||
# Non-existent currency
|
||||
latest_btc = Balance.get_latest(stake_currency="BTC")
|
||||
assert latest_btc is None
|
||||
|
||||
def test_get_latest_before(self):
|
||||
"""Test getting balance before a timestamp"""
|
||||
self._create_test_balances()
|
||||
base_time = datetime(2024, 1, 1, 12, 0, 0, tzinfo=UTC)
|
||||
before = Balance.get_latest_before(base_time + timedelta(hours=2, minutes=30))
|
||||
assert before is not None
|
||||
assert before.event_type == "trade_close"
|
||||
|
||||
def test_get_latest_before_none(self):
|
||||
"""Test getting balance before first record returns None"""
|
||||
self._create_test_balances()
|
||||
base_time = datetime(2024, 1, 1, 11, 0, 0, tzinfo=UTC)
|
||||
before = Balance.get_latest_before(base_time)
|
||||
assert before is None
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("init_persistence")
|
||||
class TestExternalChangeDetection:
|
||||
def test_calculate_external_change_no_previous(self):
|
||||
"""Test external change with no previous record"""
|
||||
balance = Balance(
|
||||
timestamp=dt_now(),
|
||||
event_type=BalanceEventType.BOT_START.value,
|
||||
total_balance=1000.0,
|
||||
free_balance=1000.0,
|
||||
stake_currency="USDT",
|
||||
total_profit=0.0,
|
||||
is_dry_run=True,
|
||||
)
|
||||
Balance.session.add(balance)
|
||||
Trade.commit()
|
||||
|
||||
external_change = Balance.calculate_external_change(balance)
|
||||
assert external_change is None
|
||||
|
||||
def test_calculate_external_change_normal_trade(self):
|
||||
"""Test external change with normal trade profit"""
|
||||
base_time = datetime(2024, 1, 1, 12, 0, 0, tzinfo=UTC)
|
||||
|
||||
# Initial balance
|
||||
b1 = Balance(
|
||||
timestamp=base_time,
|
||||
event_type=BalanceEventType.BOT_START.value,
|
||||
total_balance=1000.0,
|
||||
free_balance=1000.0,
|
||||
stake_currency="USDT",
|
||||
total_profit=0.0,
|
||||
is_dry_run=True,
|
||||
)
|
||||
Balance.session.add(b1)
|
||||
Trade.commit()
|
||||
|
||||
# Balance after profit (no external change)
|
||||
b2 = Balance(
|
||||
timestamp=base_time + timedelta(hours=1),
|
||||
event_type=BalanceEventType.TRADE_CLOSE.value,
|
||||
total_balance=1050.0, # 50 profit
|
||||
free_balance=1050.0,
|
||||
stake_currency="USDT",
|
||||
total_profit=50.0, # 50 profit
|
||||
is_dry_run=True,
|
||||
)
|
||||
|
||||
external_change = Balance.calculate_external_change(b2)
|
||||
assert external_change == 0.0 # Expected: 1000 + 50 = 1050, actual = 1050
|
||||
|
||||
def test_calculate_external_change_deposit(self):
|
||||
"""Test detecting a deposit"""
|
||||
base_time = datetime(2024, 1, 1, 12, 0, 0, tzinfo=UTC)
|
||||
|
||||
# Initial balance
|
||||
b1 = Balance(
|
||||
timestamp=base_time,
|
||||
event_type=BalanceEventType.BOT_START.value,
|
||||
total_balance=1000.0,
|
||||
free_balance=1000.0,
|
||||
stake_currency="USDT",
|
||||
total_profit=0.0,
|
||||
is_dry_run=True,
|
||||
)
|
||||
Balance.session.add(b1)
|
||||
Trade.commit()
|
||||
|
||||
# Balance after deposit
|
||||
b2 = Balance(
|
||||
timestamp=base_time + timedelta(hours=1),
|
||||
event_type=BalanceEventType.BOT_START.value,
|
||||
total_balance=1500.0, # 500 deposited
|
||||
free_balance=1500.0,
|
||||
stake_currency="USDT",
|
||||
total_profit=0.0, # No profit change
|
||||
is_dry_run=True,
|
||||
)
|
||||
|
||||
external_change = Balance.calculate_external_change(b2)
|
||||
assert external_change == 500.0 # Expected: 1000, actual: 1500
|
||||
|
||||
def test_calculate_external_change_withdrawal(self):
|
||||
"""Test detecting a withdrawal"""
|
||||
base_time = datetime(2024, 1, 1, 12, 0, 0, tzinfo=UTC)
|
||||
|
||||
# Initial balance
|
||||
b1 = Balance(
|
||||
timestamp=base_time,
|
||||
event_type=BalanceEventType.BOT_START.value,
|
||||
total_balance=1000.0,
|
||||
free_balance=1000.0,
|
||||
stake_currency="USDT",
|
||||
total_profit=0.0,
|
||||
is_dry_run=True,
|
||||
)
|
||||
Balance.session.add(b1)
|
||||
Trade.commit()
|
||||
|
||||
# Balance after withdrawal
|
||||
b2 = Balance(
|
||||
timestamp=base_time + timedelta(hours=1),
|
||||
event_type=BalanceEventType.BOT_START.value,
|
||||
total_balance=700.0, # 300 withdrawn
|
||||
free_balance=700.0,
|
||||
stake_currency="USDT",
|
||||
total_profit=0.0, # No profit change
|
||||
is_dry_run=True,
|
||||
)
|
||||
|
||||
external_change = Balance.calculate_external_change(b2)
|
||||
assert external_change == -300.0 # Expected: 1000, actual: 700
|
||||
|
||||
def test_detect_deposits_withdrawals(self):
|
||||
"""Test detecting deposits and withdrawals in batch"""
|
||||
base_time = datetime(2024, 1, 1, 12, 0, 0, tzinfo=UTC)
|
||||
|
||||
balances = [
|
||||
Balance(
|
||||
timestamp=base_time,
|
||||
event_type=BalanceEventType.BOT_START.value,
|
||||
total_balance=1000.0,
|
||||
free_balance=1000.0,
|
||||
stake_currency="USDT",
|
||||
total_profit=0.0,
|
||||
external_change=None, # First record
|
||||
is_dry_run=True,
|
||||
),
|
||||
Balance(
|
||||
timestamp=base_time + timedelta(hours=1),
|
||||
event_type=BalanceEventType.TRADE_CLOSE.value,
|
||||
total_balance=1050.0,
|
||||
free_balance=1050.0,
|
||||
stake_currency="USDT",
|
||||
total_profit=50.0,
|
||||
external_change=0.0, # No external change
|
||||
is_dry_run=True,
|
||||
),
|
||||
Balance(
|
||||
timestamp=base_time + timedelta(hours=2),
|
||||
event_type=BalanceEventType.BOT_START.value,
|
||||
total_balance=1550.0,
|
||||
free_balance=1550.0,
|
||||
stake_currency="USDT",
|
||||
total_profit=50.0,
|
||||
external_change=500.0, # Deposit
|
||||
is_dry_run=True,
|
||||
),
|
||||
Balance(
|
||||
timestamp=base_time + timedelta(hours=3),
|
||||
event_type=BalanceEventType.MANUAL.value,
|
||||
total_balance=1350.0,
|
||||
free_balance=1350.0,
|
||||
stake_currency="USDT",
|
||||
total_profit=50.0,
|
||||
external_change=-200.0, # Withdrawal
|
||||
is_dry_run=True,
|
||||
),
|
||||
]
|
||||
for b in balances:
|
||||
Balance.session.add(b)
|
||||
Trade.commit()
|
||||
|
||||
changes = Balance.detect_deposits_withdrawals()
|
||||
assert len(changes) == 2
|
||||
|
||||
# Check deposit
|
||||
deposit = next((c for c in changes if c["type"] == "deposit"), None)
|
||||
assert deposit is not None
|
||||
assert deposit["amount"] == 500.0
|
||||
|
||||
# Check withdrawal
|
||||
withdrawal = next((c for c in changes if c["type"] == "withdrawal"), None)
|
||||
assert withdrawal is not None
|
||||
assert withdrawal["amount"] == -200.0
|
||||
|
||||
def test_detect_deposits_withdrawals_with_threshold(self):
|
||||
"""Test filtering by threshold"""
|
||||
base_time = datetime(2024, 1, 1, 12, 0, 0, tzinfo=UTC)
|
||||
|
||||
balances = [
|
||||
Balance(
|
||||
timestamp=base_time,
|
||||
event_type=BalanceEventType.BOT_START.value,
|
||||
total_balance=1000.0,
|
||||
free_balance=1000.0,
|
||||
stake_currency="USDT",
|
||||
external_change=50.0, # Small deposit
|
||||
is_dry_run=True,
|
||||
),
|
||||
Balance(
|
||||
timestamp=base_time + timedelta(hours=1),
|
||||
event_type=BalanceEventType.BOT_START.value,
|
||||
total_balance=2000.0,
|
||||
free_balance=2000.0,
|
||||
stake_currency="USDT",
|
||||
external_change=500.0, # Large deposit
|
||||
is_dry_run=True,
|
||||
),
|
||||
]
|
||||
for b in balances:
|
||||
Balance.session.add(b)
|
||||
Trade.commit()
|
||||
|
||||
# With threshold of 100, should only see the large deposit
|
||||
changes = Balance.detect_deposits_withdrawals(threshold=100.0)
|
||||
assert len(changes) == 1
|
||||
assert changes[0]["amount"] == 500.0
|
||||
Loading…
Reference in new issue