diff --git a/src/database.py b/src/database.py new file mode 100644 index 0000000..cafa1ea --- /dev/null +++ b/src/database.py @@ -0,0 +1,72 @@ +import json +from datetime import datetime, timezone +from notion_client import Client +from loguru import logger + + +class TradeRepository: + """Notion 데이터베이스에 거래 이력을 저장하는 레포지토리.""" + + def __init__(self, token: str, database_id: str): + self.client = Client(auth=token) + self.database_id = database_id + + def save_trade( + self, + symbol: str, + side: str, + entry_price: float, + quantity: float, + leverage: int, + signal_data: dict = None, + ) -> dict: + properties = { + "Symbol": {"title": [{"text": {"content": symbol}}]}, + "Side": {"select": {"name": side}}, + "Entry Price": {"number": entry_price}, + "Quantity": {"number": quantity}, + "Leverage": {"number": leverage}, + "Status": {"select": {"name": "OPEN"}}, + "Signal Data": { + "rich_text": [ + {"text": {"content": json.dumps(signal_data or {}, ensure_ascii=False)}} + ] + }, + "Opened At": { + "date": {"start": datetime.now(timezone.utc).isoformat()} + }, + } + result = self.client.pages.create( + parent={"database_id": self.database_id}, + properties=properties, + ) + logger.info(f"거래 저장: {result['id']}") + return result + + def close_trade(self, trade_id: str, exit_price: float, pnl: float) -> dict: + properties = { + "Exit Price": {"number": exit_price}, + "PnL": {"number": pnl}, + "Status": {"select": {"name": "CLOSED"}}, + "Closed At": { + "date": {"start": datetime.now(timezone.utc).isoformat()} + }, + } + result = self.client.pages.update( + page_id=trade_id, + properties=properties, + ) + logger.info(f"거래 종료: {trade_id}, PnL: {pnl:.4f}") + return result + + def get_open_trades(self, symbol: str) -> list[dict]: + response = self.client.databases.query( + database_id=self.database_id, + filter={ + "and": [ + {"property": "Symbol", "title": {"equals": symbol}}, + {"property": "Status", "select": {"equals": "OPEN"}}, + ] + }, + ) + return response.get("results", []) diff --git a/tests/test_database.py b/tests/test_database.py new file mode 100644 index 0000000..1b29d90 --- /dev/null +++ b/tests/test_database.py @@ -0,0 +1,42 @@ +import pytest +from unittest.mock import MagicMock, patch +from src.database import TradeRepository + + +@pytest.fixture +def mock_repo(): + with patch("src.database.Client") as mock_client_cls: + mock_client = MagicMock() + mock_client_cls.return_value = mock_client + repo = TradeRepository(token="secret_test", database_id="db_test") + repo.client = mock_client + yield repo + + +def test_save_trade(mock_repo): + mock_repo.client.pages.create.return_value = { + "id": "abc123", + "properties": {}, + } + result = mock_repo.save_trade( + symbol="XRPUSDT", + side="LONG", + entry_price=0.5, + quantity=400.0, + leverage=10, + signal_data={"rsi": 32, "macd_hist": 0.001}, + ) + assert result["id"] == "abc123" + + +def test_close_trade(mock_repo): + mock_repo.client.pages.update.return_value = { + "id": "abc123", + "properties": { + "Status": {"select": {"name": "CLOSED"}}, + }, + } + result = mock_repo.close_trade( + trade_id="abc123", exit_price=0.55, pnl=20.0 + ) + assert result["id"] == "abc123"