diff --git a/alembic.ini b/alembic.ini new file mode 100644 index 0000000..b98ba1a --- /dev/null +++ b/alembic.ini @@ -0,0 +1,79 @@ +# Alembic Configuration File + +[alembic] +# Path to migration scripts +script_location = alembic + +# Template used to generate migration files +file_template = %%(year)d%%(month).2d%%(day).2d_%%(hour).2d%%(minute).2d_%%(slug)s + +# Prepend timestamp to migration file names +prepend_date = True + +# Timezone to use when rendering the date within the migration file +# as well as the filename. Use UTC for consistency. +timezone = UTC + +# Max length of characters to apply to the "slug" field +truncate_slug_length = 40 + +# Set to 'true' to run the environment during the 'revision' command +revision_environment = false + +# Set to 'true' to allow .pyc and .pyo files without a source .py file +sourceless = false + +# Version location specification +version_locations = %(here)s/alembic/versions + +# Version path separator +version_path_separator = os + +# Database URL - override with SQLALCHEMY_DATABASE_URL environment variable +sqlalchemy.url = postgresql://localhost/polymarket_tracker + + +[post_write_hooks] +# Post write hooks define scripts to run after generating new revision files + +# Format using "black" - only if available +# hooks = black +# black.type = console_scripts +# black.entrypoint = black +# black.options = -q + + +# Logging configuration +[loggers] +keys = root,sqlalchemy,alembic + +[handlers] +keys = console + +[formatters] +keys = generic + +[logger_root] +level = WARN +handlers = console +qualname = + +[logger_sqlalchemy] +level = WARN +handlers = +qualname = sqlalchemy.engine + +[logger_alembic] +level = INFO +handlers = +qualname = alembic + +[handler_console] +class = StreamHandler +args = (sys.stderr,) +level = NOTSET +formatter = generic + +[formatter_generic] +format = %(levelname)-5.5s [%(name)s] %(message)s +datefmt = %H:%M:%S diff --git a/alembic/env.py b/alembic/env.py new file mode 100644 index 0000000..27d4062 --- /dev/null +++ b/alembic/env.py @@ -0,0 +1,71 @@ +"""Alembic migration environment configuration.""" + +import os +from logging.config import fileConfig + +from alembic import context +from sqlalchemy import engine_from_config, pool + +from polymarket_insider_tracker.storage.models import Base + +# this is the Alembic Config object +config = context.config + +# Interpret the config file for Python logging +if config.config_file_name is not None: + fileConfig(config.config_file_name) + +# Target metadata for 'autogenerate' support +target_metadata = Base.metadata + +# Get database URL from environment variable or config +database_url = os.environ.get("SQLALCHEMY_DATABASE_URL") +if database_url: + config.set_main_option("sqlalchemy.url", database_url) + + +def run_migrations_offline() -> None: + """Run migrations in 'offline' mode. + + This configures the context with just a URL and not an Engine, + though an Engine is acceptable here as well. By skipping the Engine + creation we don't even need a DBAPI to be available. + + Calls to context.execute() here emit the given string to the + script output. + """ + url = config.get_main_option("sqlalchemy.url") + context.configure( + url=url, + target_metadata=target_metadata, + literal_binds=True, + dialect_opts={"paramstyle": "named"}, + ) + + with context.begin_transaction(): + context.run_migrations() + + +def run_migrations_online() -> None: + """Run migrations in 'online' mode. + + In this scenario we need to create an Engine and associate a + connection with the context. + """ + connectable = engine_from_config( + config.get_section(config.config_ini_section, {}), + prefix="sqlalchemy.", + poolclass=pool.NullPool, + ) + + with connectable.connect() as connection: + context.configure(connection=connection, target_metadata=target_metadata) + + with context.begin_transaction(): + context.run_migrations() + + +if context.is_offline_mode(): + run_migrations_offline() +else: + run_migrations_online() diff --git a/alembic/script.py.mako b/alembic/script.py.mako new file mode 100644 index 0000000..fbc4b07 --- /dev/null +++ b/alembic/script.py.mako @@ -0,0 +1,26 @@ +"""${message} + +Revision ID: ${up_revision} +Revises: ${down_revision | comma,n} +Create Date: ${create_date} + +""" +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa +${imports if imports else ""} + +# revision identifiers, used by Alembic. +revision: str = ${repr(up_revision)} +down_revision: Union[str, None] = ${repr(down_revision)} +branch_labels: Union[str, Sequence[str], None] = ${repr(branch_labels)} +depends_on: Union[str, Sequence[str], None] = ${repr(depends_on)} + + +def upgrade() -> None: + ${upgrades if upgrades else "pass"} + + +def downgrade() -> None: + ${downgrades if downgrades else "pass"} diff --git a/alembic/versions/20260104_0000_initial_schema.py b/alembic/versions/20260104_0000_initial_schema.py new file mode 100644 index 0000000..9ab69b6 --- /dev/null +++ b/alembic/versions/20260104_0000_initial_schema.py @@ -0,0 +1,87 @@ +"""Initial schema for wallet profiles and funding transfers. + +Revision ID: 001_initial +Revises: +Create Date: 2026-01-04 00:00:00.000000+00:00 +""" + +from typing import Sequence, Union + +import sqlalchemy as sa +from alembic import op + +# revision identifiers, used by Alembic. +revision: str = "001_initial" +down_revision: Union[str, None] = None +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + # Wallet profiles table + op.create_table( + "wallet_profiles", + sa.Column("id", sa.Integer(), autoincrement=True, nullable=False), + sa.Column("address", sa.String(42), nullable=False), + sa.Column("nonce", sa.Integer(), nullable=False), + sa.Column("first_seen_at", sa.DateTime(timezone=True), nullable=True), + sa.Column("is_fresh", sa.Boolean(), nullable=False), + sa.Column("matic_balance", sa.Numeric(30, 0), nullable=True), + sa.Column("usdc_balance", sa.Numeric(20, 6), nullable=True), + sa.Column("analyzed_at", sa.DateTime(timezone=True), nullable=False), + sa.Column("created_at", sa.DateTime(timezone=True), nullable=False), + sa.Column("updated_at", sa.DateTime(timezone=True), nullable=False), + sa.PrimaryKeyConstraint("id"), + sa.UniqueConstraint("address"), + ) + op.create_index("idx_wallet_profiles_address", "wallet_profiles", ["address"]) + + # Funding transfers table + op.create_table( + "funding_transfers", + sa.Column("id", sa.Integer(), autoincrement=True, nullable=False), + sa.Column("from_address", sa.String(42), nullable=False), + sa.Column("to_address", sa.String(42), nullable=False), + sa.Column("amount", sa.Numeric(30, 6), nullable=False), + sa.Column("token", sa.String(10), nullable=False), + sa.Column("tx_hash", sa.String(66), nullable=False), + sa.Column("block_number", sa.Integer(), nullable=False), + sa.Column("timestamp", sa.DateTime(timezone=True), nullable=False), + sa.Column("created_at", sa.DateTime(timezone=True), nullable=False), + sa.PrimaryKeyConstraint("id"), + sa.UniqueConstraint("tx_hash"), + ) + op.create_index("idx_funding_transfers_to", "funding_transfers", ["to_address"]) + op.create_index("idx_funding_transfers_from", "funding_transfers", ["from_address"]) + op.create_index("idx_funding_transfers_block", "funding_transfers", ["block_number"]) + + # Wallet relationships table + op.create_table( + "wallet_relationships", + sa.Column("id", sa.Integer(), autoincrement=True, nullable=False), + sa.Column("wallet_a", sa.String(42), nullable=False), + sa.Column("wallet_b", sa.String(42), nullable=False), + sa.Column("relationship_type", sa.String(20), nullable=False), + sa.Column("confidence", sa.Numeric(3, 2), nullable=False), + sa.Column("created_at", sa.DateTime(timezone=True), nullable=False), + sa.PrimaryKeyConstraint("id"), + sa.UniqueConstraint( + "wallet_a", "wallet_b", "relationship_type", name="uq_wallet_relationship" + ), + ) + op.create_index("idx_wallet_relationships_a", "wallet_relationships", ["wallet_a"]) + op.create_index("idx_wallet_relationships_b", "wallet_relationships", ["wallet_b"]) + + +def downgrade() -> None: + op.drop_index("idx_wallet_relationships_b", table_name="wallet_relationships") + op.drop_index("idx_wallet_relationships_a", table_name="wallet_relationships") + op.drop_table("wallet_relationships") + + op.drop_index("idx_funding_transfers_block", table_name="funding_transfers") + op.drop_index("idx_funding_transfers_from", table_name="funding_transfers") + op.drop_index("idx_funding_transfers_to", table_name="funding_transfers") + op.drop_table("funding_transfers") + + op.drop_index("idx_wallet_profiles_address", table_name="wallet_profiles") + op.drop_table("wallet_profiles") diff --git a/src/polymarket_insider_tracker/storage/__init__.py b/src/polymarket_insider_tracker/storage/__init__.py index 41ea10c..6f75660 100644 --- a/src/polymarket_insider_tracker/storage/__init__.py +++ b/src/polymarket_insider_tracker/storage/__init__.py @@ -1 +1,45 @@ """Storage layer - Database schemas and repositories.""" + +from polymarket_insider_tracker.storage.database import ( + DatabaseManager, + create_async_db_engine, + create_async_session_factory, + create_sync_engine, + create_sync_session_factory, + init_async_db, + init_db, +) +from polymarket_insider_tracker.storage.models import ( + Base, + FundingTransferModel, + WalletProfileModel, + WalletRelationshipModel, +) +from polymarket_insider_tracker.storage.repos import ( + FundingRepository, + FundingTransferDTO, + RelationshipRepository, + WalletProfileDTO, + WalletRelationshipDTO, + WalletRepository, +) + +__all__ = [ + "Base", + "DatabaseManager", + "FundingRepository", + "FundingTransferDTO", + "FundingTransferModel", + "RelationshipRepository", + "WalletProfileDTO", + "WalletProfileModel", + "WalletRelationshipDTO", + "WalletRelationshipModel", + "WalletRepository", + "create_async_db_engine", + "create_async_session_factory", + "create_sync_engine", + "create_sync_session_factory", + "init_async_db", + "init_db", +] diff --git a/src/polymarket_insider_tracker/storage/database.py b/src/polymarket_insider_tracker/storage/database.py new file mode 100644 index 0000000..4c847ea --- /dev/null +++ b/src/polymarket_insider_tracker/storage/database.py @@ -0,0 +1,209 @@ +"""Database connection and session management. + +This module provides the database engine, session factory, and +async session support for the storage layer. +""" + +from __future__ import annotations + +import logging +from collections.abc import AsyncGenerator +from contextlib import asynccontextmanager +from typing import TYPE_CHECKING, Any + +from sqlalchemy import create_engine +from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine +from sqlalchemy.orm import Session, sessionmaker + +from polymarket_insider_tracker.storage.models import Base + +if TYPE_CHECKING: + from sqlalchemy import Engine + from sqlalchemy.ext.asyncio import AsyncEngine + +logger = logging.getLogger(__name__) + + +def create_sync_engine(database_url: str, **kwargs: Any) -> Engine: + """Create a synchronous SQLAlchemy engine. + + Args: + database_url: Database connection URL (e.g., postgresql://...). + **kwargs: Additional engine options. + + Returns: + SQLAlchemy Engine instance. + """ + return create_engine(database_url, **kwargs) + + +def create_async_db_engine(database_url: str, **kwargs: Any) -> AsyncEngine: + """Create an asynchronous SQLAlchemy engine. + + Args: + database_url: Database connection URL (e.g., postgresql+asyncpg://...). + **kwargs: Additional engine options. + + Returns: + SQLAlchemy AsyncEngine instance. + """ + return create_async_engine(database_url, **kwargs) + + +def create_sync_session_factory(engine: Engine) -> sessionmaker[Session]: + """Create a synchronous session factory. + + Args: + engine: SQLAlchemy Engine instance. + + Returns: + Session factory. + """ + return sessionmaker(bind=engine, expire_on_commit=False) + + +def create_async_session_factory(engine: AsyncEngine) -> async_sessionmaker[AsyncSession]: + """Create an asynchronous session factory. + + Args: + engine: SQLAlchemy AsyncEngine instance. + + Returns: + Async session factory. + """ + return async_sessionmaker(bind=engine, expire_on_commit=False) + + +def init_db(engine: Engine) -> None: + """Initialize the database schema. + + Creates all tables defined in the models. + + Args: + engine: SQLAlchemy Engine instance. + """ + Base.metadata.create_all(engine) + logger.info("Database schema initialized") + + +async def init_async_db(engine: AsyncEngine) -> None: + """Initialize the database schema asynchronously. + + Creates all tables defined in the models. + + Args: + engine: SQLAlchemy AsyncEngine instance. + """ + async with engine.begin() as conn: + await conn.run_sync(Base.metadata.create_all) + logger.info("Database schema initialized (async)") + + +class DatabaseManager: + """Manages database connections and sessions. + + Provides a unified interface for both sync and async database operations. + """ + + def __init__( + self, + database_url: str, + *, + async_mode: bool = True, + pool_size: int = 5, + max_overflow: int = 10, + echo: bool = False, + ) -> None: + """Initialize database manager. + + Args: + database_url: Database connection URL. + async_mode: Use async engine/sessions if True. + pool_size: Connection pool size. + max_overflow: Maximum overflow connections. + echo: Echo SQL statements for debugging. + """ + self.database_url = database_url + self.async_mode = async_mode + self._pool_size = pool_size + self._max_overflow = max_overflow + self._echo = echo + + self._sync_engine: Engine | None = None + self._async_engine: AsyncEngine | None = None + self._sync_session_factory: sessionmaker[Session] | None = None + self._async_session_factory: async_sessionmaker[AsyncSession] | None = None + + def _get_sync_engine(self) -> Engine: + """Get or create the synchronous engine.""" + if self._sync_engine is None: + self._sync_engine = create_sync_engine( + self.database_url, + pool_size=self._pool_size, + max_overflow=self._max_overflow, + echo=self._echo, + ) + return self._sync_engine + + def _get_async_engine(self) -> AsyncEngine: + """Get or create the asynchronous engine.""" + if self._async_engine is None: + self._async_engine = create_async_db_engine( + self.database_url, + pool_size=self._pool_size, + max_overflow=self._max_overflow, + echo=self._echo, + ) + return self._async_engine + + def get_sync_session(self) -> Session: + """Get a new synchronous session. + + Returns: + SQLAlchemy Session instance. + """ + if self._sync_session_factory is None: + self._sync_session_factory = create_sync_session_factory(self._get_sync_engine()) + return self._sync_session_factory() + + @asynccontextmanager + async def get_async_session(self) -> AsyncGenerator[AsyncSession, None]: + """Get an asynchronous session as a context manager. + + Yields: + SQLAlchemy AsyncSession instance. + """ + if self._async_session_factory is None: + self._async_session_factory = create_async_session_factory(self._get_async_engine()) + + session = self._async_session_factory() + try: + yield session + await session.commit() + except Exception: + await session.rollback() + raise + finally: + await session.close() + + def init_schema(self) -> None: + """Initialize database schema synchronously.""" + init_db(self._get_sync_engine()) + + async def init_schema_async(self) -> None: + """Initialize database schema asynchronously.""" + await init_async_db(self._get_async_engine()) + + def dispose(self) -> None: + """Dispose of all database connections.""" + if self._sync_engine is not None: + self._sync_engine.dispose() + self._sync_engine = None + logger.info("Database connections disposed") + + async def dispose_async(self) -> None: + """Dispose of all async database connections.""" + if self._async_engine is not None: + await self._async_engine.dispose() + self._async_engine = None + logger.info("Async database connections disposed") diff --git a/src/polymarket_insider_tracker/storage/models.py b/src/polymarket_insider_tracker/storage/models.py new file mode 100644 index 0000000..ace3058 --- /dev/null +++ b/src/polymarket_insider_tracker/storage/models.py @@ -0,0 +1,115 @@ +"""SQLAlchemy models for persistent storage. + +This module defines the database schema for storing wallet profiles, +funding transfers, and wallet relationships. +""" + +from __future__ import annotations + +from datetime import UTC, datetime +from decimal import Decimal +from typing import TYPE_CHECKING + +from sqlalchemy import ( + Boolean, + DateTime, + Index, + Integer, + Numeric, + String, + UniqueConstraint, +) +from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column + +if TYPE_CHECKING: + pass + + +class Base(DeclarativeBase): + """Base class for all SQLAlchemy models.""" + + pass + + +class WalletProfileModel(Base): + """SQLAlchemy model for wallet profiles. + + Stores analyzed wallet information including age, transaction count, + balances, and freshness classification. + """ + + __tablename__ = "wallet_profiles" + + id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True) + address: Mapped[str] = mapped_column(String(42), unique=True, nullable=False) + nonce: Mapped[int] = mapped_column(Integer, nullable=False) + first_seen_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True) + is_fresh: Mapped[bool] = mapped_column(Boolean, nullable=False) + matic_balance: Mapped[Decimal | None] = mapped_column(Numeric(30, 0), nullable=True) + usdc_balance: Mapped[Decimal | None] = mapped_column(Numeric(20, 6), nullable=True) + analyzed_at: Mapped[datetime] = mapped_column( + DateTime(timezone=True), nullable=False, default=lambda: datetime.now(UTC) + ) + created_at: Mapped[datetime] = mapped_column( + DateTime(timezone=True), nullable=False, default=lambda: datetime.now(UTC) + ) + updated_at: Mapped[datetime] = mapped_column( + DateTime(timezone=True), + nullable=False, + default=lambda: datetime.now(UTC), + onupdate=lambda: datetime.now(UTC), + ) + + __table_args__ = (Index("idx_wallet_profiles_address", "address"),) + + +class FundingTransferModel(Base): + """SQLAlchemy model for funding transfers. + + Stores ERC20 transfer events to track wallet funding sources. + """ + + __tablename__ = "funding_transfers" + + id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True) + from_address: Mapped[str] = mapped_column(String(42), nullable=False) + to_address: Mapped[str] = mapped_column(String(42), nullable=False) + amount: Mapped[Decimal] = mapped_column(Numeric(30, 6), nullable=False) + token: Mapped[str] = mapped_column(String(10), nullable=False) + tx_hash: Mapped[str] = mapped_column(String(66), unique=True, nullable=False) + block_number: Mapped[int] = mapped_column(Integer, nullable=False) + timestamp: Mapped[datetime] = mapped_column(DateTime(timezone=True), nullable=False) + created_at: Mapped[datetime] = mapped_column( + DateTime(timezone=True), nullable=False, default=lambda: datetime.now(UTC) + ) + + __table_args__ = ( + Index("idx_funding_transfers_to", "to_address"), + Index("idx_funding_transfers_from", "from_address"), + Index("idx_funding_transfers_block", "block_number"), + ) + + +class WalletRelationshipModel(Base): + """SQLAlchemy model for wallet relationships. + + Stores graph edges between wallets representing funding relationships + or entity linkages. + """ + + __tablename__ = "wallet_relationships" + + id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True) + wallet_a: Mapped[str] = mapped_column(String(42), nullable=False) + wallet_b: Mapped[str] = mapped_column(String(42), nullable=False) + relationship_type: Mapped[str] = mapped_column(String(20), nullable=False) + confidence: Mapped[Decimal] = mapped_column(Numeric(3, 2), nullable=False) + created_at: Mapped[datetime] = mapped_column( + DateTime(timezone=True), nullable=False, default=lambda: datetime.now(UTC) + ) + + __table_args__ = ( + UniqueConstraint("wallet_a", "wallet_b", "relationship_type", name="uq_wallet_relationship"), + Index("idx_wallet_relationships_a", "wallet_a"), + Index("idx_wallet_relationships_b", "wallet_b"), + ) diff --git a/src/polymarket_insider_tracker/storage/repos.py b/src/polymarket_insider_tracker/storage/repos.py new file mode 100644 index 0000000..3a5c447 --- /dev/null +++ b/src/polymarket_insider_tracker/storage/repos.py @@ -0,0 +1,515 @@ +"""Repository pattern implementations for data access. + +This module provides clean data access abstractions for wallet profiles, +funding transfers, and wallet relationships. +""" + +from __future__ import annotations + +import logging +from dataclasses import dataclass +from datetime import UTC, datetime +from decimal import Decimal +from typing import TYPE_CHECKING + +from sqlalchemy import delete, select, update +from sqlalchemy.dialects.postgresql import insert as pg_insert +from sqlalchemy.dialects.sqlite import insert as sqlite_insert + +from polymarket_insider_tracker.storage.models import ( + FundingTransferModel, + WalletProfileModel, + WalletRelationshipModel, +) + +if TYPE_CHECKING: + from sqlalchemy.ext.asyncio import AsyncSession + +logger = logging.getLogger(__name__) + + +@dataclass +class WalletProfileDTO: + """Data transfer object for wallet profiles.""" + + address: str + nonce: int + first_seen_at: datetime | None + is_fresh: bool + matic_balance: Decimal | None + usdc_balance: Decimal | None + analyzed_at: datetime + created_at: datetime | None = None + updated_at: datetime | None = None + + @classmethod + def from_model(cls, model: WalletProfileModel) -> WalletProfileDTO: + """Create DTO from SQLAlchemy model.""" + return cls( + address=model.address, + nonce=model.nonce, + first_seen_at=model.first_seen_at, + is_fresh=model.is_fresh, + matic_balance=model.matic_balance, + usdc_balance=model.usdc_balance, + analyzed_at=model.analyzed_at, + created_at=model.created_at, + updated_at=model.updated_at, + ) + + +@dataclass +class FundingTransferDTO: + """Data transfer object for funding transfers.""" + + from_address: str + to_address: str + amount: Decimal + token: str + tx_hash: str + block_number: int + timestamp: datetime + created_at: datetime | None = None + + @classmethod + def from_model(cls, model: FundingTransferModel) -> FundingTransferDTO: + """Create DTO from SQLAlchemy model.""" + return cls( + from_address=model.from_address, + to_address=model.to_address, + amount=model.amount, + token=model.token, + tx_hash=model.tx_hash, + block_number=model.block_number, + timestamp=model.timestamp, + created_at=model.created_at, + ) + + +@dataclass +class WalletRelationshipDTO: + """Data transfer object for wallet relationships.""" + + wallet_a: str + wallet_b: str + relationship_type: str + confidence: Decimal + created_at: datetime | None = None + + @classmethod + def from_model(cls, model: WalletRelationshipModel) -> WalletRelationshipDTO: + """Create DTO from SQLAlchemy model.""" + return cls( + wallet_a=model.wallet_a, + wallet_b=model.wallet_b, + relationship_type=model.relationship_type, + confidence=model.confidence, + created_at=model.created_at, + ) + + +class WalletRepository: + """Repository for wallet profile data access. + + Provides CRUD operations for wallet profiles with async support. + """ + + def __init__(self, session: AsyncSession) -> None: + """Initialize repository with database session. + + Args: + session: SQLAlchemy async session. + """ + self.session = session + + async def get_by_address(self, address: str) -> WalletProfileDTO | None: + """Get wallet profile by address. + + Args: + address: Wallet address (lowercase). + + Returns: + WalletProfileDTO if found, None otherwise. + """ + result = await self.session.execute( + select(WalletProfileModel).where(WalletProfileModel.address == address.lower()) + ) + model = result.scalar_one_or_none() + return WalletProfileDTO.from_model(model) if model else None + + async def get_many(self, addresses: list[str]) -> list[WalletProfileDTO]: + """Get multiple wallet profiles by addresses. + + Args: + addresses: List of wallet addresses. + + Returns: + List of WalletProfileDTOs for found addresses. + """ + normalized = [addr.lower() for addr in addresses] + result = await self.session.execute( + select(WalletProfileModel).where(WalletProfileModel.address.in_(normalized)) + ) + return [WalletProfileDTO.from_model(m) for m in result.scalars().all()] + + async def get_fresh_wallets(self, limit: int = 100) -> list[WalletProfileDTO]: + """Get recent fresh wallets. + + Args: + limit: Maximum number of results. + + Returns: + List of WalletProfileDTOs marked as fresh. + """ + result = await self.session.execute( + select(WalletProfileModel) + .where(WalletProfileModel.is_fresh.is_(True)) + .order_by(WalletProfileModel.analyzed_at.desc()) + .limit(limit) + ) + return [WalletProfileDTO.from_model(m) for m in result.scalars().all()] + + async def upsert(self, dto: WalletProfileDTO) -> WalletProfileDTO: + """Insert or update wallet profile. + + Args: + dto: Wallet profile data. + + Returns: + Updated WalletProfileDTO. + """ + now = datetime.now(UTC) + values = { + "address": dto.address.lower(), + "nonce": dto.nonce, + "first_seen_at": dto.first_seen_at, + "is_fresh": dto.is_fresh, + "matic_balance": dto.matic_balance, + "usdc_balance": dto.usdc_balance, + "analyzed_at": dto.analyzed_at, + "updated_at": now, + } + + # Try PostgreSQL upsert first, fall back to SQLite for testing + try: + stmt = pg_insert(WalletProfileModel).values(**values, created_at=now) + stmt = stmt.on_conflict_do_update( + index_elements=["address"], + set_={ + "nonce": stmt.excluded.nonce, + "first_seen_at": stmt.excluded.first_seen_at, + "is_fresh": stmt.excluded.is_fresh, + "matic_balance": stmt.excluded.matic_balance, + "usdc_balance": stmt.excluded.usdc_balance, + "analyzed_at": stmt.excluded.analyzed_at, + "updated_at": stmt.excluded.updated_at, + }, + ) + await self.session.execute(stmt) + except Exception: + # Fall back to SQLite upsert for testing + stmt = sqlite_insert(WalletProfileModel).values(**values, created_at=now) + stmt = stmt.on_conflict_do_update( + index_elements=["address"], + set_={ + "nonce": stmt.excluded.nonce, + "first_seen_at": stmt.excluded.first_seen_at, + "is_fresh": stmt.excluded.is_fresh, + "matic_balance": stmt.excluded.matic_balance, + "usdc_balance": stmt.excluded.usdc_balance, + "analyzed_at": stmt.excluded.analyzed_at, + "updated_at": stmt.excluded.updated_at, + }, + ) + await self.session.execute(stmt) + + await self.session.flush() + return dto + + async def delete(self, address: str) -> bool: + """Delete wallet profile by address. + + Args: + address: Wallet address. + + Returns: + True if deleted, False if not found. + """ + result = await self.session.execute( + delete(WalletProfileModel).where(WalletProfileModel.address == address.lower()) + ) + return result.rowcount > 0 + + async def mark_stale(self, address: str) -> bool: + """Mark a wallet profile as stale (soft delete). + + Sets analyzed_at to a very old date to trigger re-analysis. + + Args: + address: Wallet address. + + Returns: + True if updated, False if not found. + """ + stale_time = datetime(2000, 1, 1, tzinfo=UTC) + result = await self.session.execute( + update(WalletProfileModel) + .where(WalletProfileModel.address == address.lower()) + .values(analyzed_at=stale_time, updated_at=datetime.now(UTC)) + ) + return result.rowcount > 0 + + +class FundingRepository: + """Repository for funding transfer data access. + + Provides CRUD operations for funding transfers with async support. + """ + + def __init__(self, session: AsyncSession) -> None: + """Initialize repository with database session. + + Args: + session: SQLAlchemy async session. + """ + self.session = session + + async def get_transfers_to( + self, address: str, limit: int = 100 + ) -> list[FundingTransferDTO]: + """Get transfers to a wallet address. + + Args: + address: Destination wallet address. + limit: Maximum number of results. + + Returns: + List of FundingTransferDTOs ordered by timestamp. + """ + result = await self.session.execute( + select(FundingTransferModel) + .where(FundingTransferModel.to_address == address.lower()) + .order_by(FundingTransferModel.timestamp.asc()) + .limit(limit) + ) + return [FundingTransferDTO.from_model(m) for m in result.scalars().all()] + + async def get_transfers_from( + self, address: str, limit: int = 100 + ) -> list[FundingTransferDTO]: + """Get transfers from a wallet address. + + Args: + address: Source wallet address. + limit: Maximum number of results. + + Returns: + List of FundingTransferDTOs ordered by timestamp. + """ + result = await self.session.execute( + select(FundingTransferModel) + .where(FundingTransferModel.from_address == address.lower()) + .order_by(FundingTransferModel.timestamp.asc()) + .limit(limit) + ) + return [FundingTransferDTO.from_model(m) for m in result.scalars().all()] + + async def get_first_transfer_to(self, address: str) -> FundingTransferDTO | None: + """Get the first transfer to a wallet. + + Args: + address: Wallet address. + + Returns: + First FundingTransferDTO if found, None otherwise. + """ + result = await self.session.execute( + select(FundingTransferModel) + .where(FundingTransferModel.to_address == address.lower()) + .order_by(FundingTransferModel.timestamp.asc()) + .limit(1) + ) + model = result.scalar_one_or_none() + return FundingTransferDTO.from_model(model) if model else None + + async def get_by_tx_hash(self, tx_hash: str) -> FundingTransferDTO | None: + """Get transfer by transaction hash. + + Args: + tx_hash: Transaction hash. + + Returns: + FundingTransferDTO if found, None otherwise. + """ + result = await self.session.execute( + select(FundingTransferModel).where(FundingTransferModel.tx_hash == tx_hash.lower()) + ) + model = result.scalar_one_or_none() + return FundingTransferDTO.from_model(model) if model else None + + async def insert(self, dto: FundingTransferDTO) -> FundingTransferDTO: + """Insert a new funding transfer. + + Args: + dto: Funding transfer data. + + Returns: + Inserted FundingTransferDTO. + + Raises: + IntegrityError if tx_hash already exists. + """ + model = FundingTransferModel( + from_address=dto.from_address.lower(), + to_address=dto.to_address.lower(), + amount=dto.amount, + token=dto.token, + tx_hash=dto.tx_hash.lower(), + block_number=dto.block_number, + timestamp=dto.timestamp, + ) + self.session.add(model) + await self.session.flush() + return dto + + async def insert_many(self, dtos: list[FundingTransferDTO]) -> int: + """Insert multiple funding transfers. + + Skips duplicates silently. + + Args: + dtos: List of funding transfer data. + + Returns: + Number of transfers inserted. + """ + inserted = 0 + for dto in dtos: + try: + await self.insert(dto) + inserted += 1 + except Exception as e: + # Skip duplicates + if "UNIQUE constraint" in str(e) or "duplicate key" in str(e).lower(): + continue + raise + return inserted + + +class RelationshipRepository: + """Repository for wallet relationship data access. + + Provides CRUD operations for wallet relationships with async support. + """ + + def __init__(self, session: AsyncSession) -> None: + """Initialize repository with database session. + + Args: + session: SQLAlchemy async session. + """ + self.session = session + + async def get_relationships( + self, wallet: str, relationship_type: str | None = None + ) -> list[WalletRelationshipDTO]: + """Get relationships for a wallet. + + Args: + wallet: Wallet address. + relationship_type: Optional filter by type. + + Returns: + List of WalletRelationshipDTOs. + """ + stmt = select(WalletRelationshipModel).where( + (WalletRelationshipModel.wallet_a == wallet.lower()) + | (WalletRelationshipModel.wallet_b == wallet.lower()) + ) + if relationship_type: + stmt = stmt.where(WalletRelationshipModel.relationship_type == relationship_type) + + result = await self.session.execute(stmt) + return [WalletRelationshipDTO.from_model(m) for m in result.scalars().all()] + + async def get_related_wallets( + self, wallet: str, relationship_type: str | None = None + ) -> list[str]: + """Get addresses of related wallets. + + Args: + wallet: Wallet address. + relationship_type: Optional filter by type. + + Returns: + List of related wallet addresses. + """ + relationships = await self.get_relationships(wallet, relationship_type) + related = set() + normalized = wallet.lower() + for rel in relationships: + if rel.wallet_a == normalized: + related.add(rel.wallet_b) + else: + related.add(rel.wallet_a) + return list(related) + + async def upsert(self, dto: WalletRelationshipDTO) -> WalletRelationshipDTO: + """Insert or update wallet relationship. + + Args: + dto: Wallet relationship data. + + Returns: + Updated WalletRelationshipDTO. + """ + now = datetime.now(UTC) + values = { + "wallet_a": dto.wallet_a.lower(), + "wallet_b": dto.wallet_b.lower(), + "relationship_type": dto.relationship_type, + "confidence": dto.confidence, + "created_at": now, + } + + # Try PostgreSQL upsert first, fall back to SQLite for testing + try: + stmt = pg_insert(WalletRelationshipModel).values(**values) + stmt = stmt.on_conflict_do_update( + constraint="uq_wallet_relationship", + set_={"confidence": stmt.excluded.confidence}, + ) + await self.session.execute(stmt) + except Exception: + # Fall back to SQLite upsert for testing + stmt = sqlite_insert(WalletRelationshipModel).values(**values) + stmt = stmt.on_conflict_do_update( + index_elements=["wallet_a", "wallet_b", "relationship_type"], + set_={"confidence": stmt.excluded.confidence}, + ) + await self.session.execute(stmt) + + await self.session.flush() + return dto + + async def delete( + self, wallet_a: str, wallet_b: str, relationship_type: str + ) -> bool: + """Delete a specific relationship. + + Args: + wallet_a: First wallet address. + wallet_b: Second wallet address. + relationship_type: Type of relationship. + + Returns: + True if deleted, False if not found. + """ + result = await self.session.execute( + delete(WalletRelationshipModel).where( + WalletRelationshipModel.wallet_a == wallet_a.lower(), + WalletRelationshipModel.wallet_b == wallet_b.lower(), + WalletRelationshipModel.relationship_type == relationship_type, + ) + ) + return result.rowcount > 0 diff --git a/tests/storage/__init__.py b/tests/storage/__init__.py index 9785e36..2725447 100644 --- a/tests/storage/__init__.py +++ b/tests/storage/__init__.py @@ -1 +1 @@ -"""Tests for storage module.""" +"""Tests for the storage layer.""" diff --git a/tests/storage/test_repos.py b/tests/storage/test_repos.py new file mode 100644 index 0000000..a31b044 --- /dev/null +++ b/tests/storage/test_repos.py @@ -0,0 +1,496 @@ +"""Tests for storage repositories.""" + +from datetime import UTC, datetime, timedelta +from decimal import Decimal + +import pytest +from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine + +from polymarket_insider_tracker.storage.models import Base +from polymarket_insider_tracker.storage.repos import ( + FundingRepository, + FundingTransferDTO, + RelationshipRepository, + WalletProfileDTO, + WalletRelationshipDTO, + WalletRepository, +) + +# ============================================================================ +# Fixtures +# ============================================================================ + + +@pytest.fixture +async def async_engine(): + """Create an async SQLite engine for testing.""" + engine = create_async_engine( + "sqlite+aiosqlite:///:memory:", + echo=False, + ) + async with engine.begin() as conn: + await conn.run_sync(Base.metadata.create_all) + yield engine + await engine.dispose() + + +@pytest.fixture +async def async_session(async_engine) -> AsyncSession: + """Create an async session for testing.""" + session_factory = async_sessionmaker(bind=async_engine, expire_on_commit=False) + async with session_factory() as session: + yield session + + +@pytest.fixture +def sample_wallet_dto() -> WalletProfileDTO: + """Create a sample wallet profile DTO.""" + return WalletProfileDTO( + address="0x1234567890abcdef1234567890abcdef12345678", + nonce=5, + first_seen_at=datetime.now(UTC) - timedelta(hours=24), + is_fresh=True, + matic_balance=Decimal("1000000000000000000"), + usdc_balance=Decimal("1000.00"), + analyzed_at=datetime.now(UTC), + ) + + +@pytest.fixture +def sample_transfer_dto() -> FundingTransferDTO: + """Create a sample funding transfer DTO.""" + return FundingTransferDTO( + from_address="0xaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", + to_address="0x1234567890abcdef1234567890abcdef12345678", + amount=Decimal("5000.00"), + token="USDC", + tx_hash="0x" + "a" * 64, + block_number=12345678, + timestamp=datetime.now(UTC), + ) + + +@pytest.fixture +def sample_relationship_dto() -> WalletRelationshipDTO: + """Create a sample wallet relationship DTO.""" + return WalletRelationshipDTO( + wallet_a="0x1234567890abcdef1234567890abcdef12345678", + wallet_b="0xaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", + relationship_type="funded_by", + confidence=Decimal("0.95"), + ) + + +# ============================================================================ +# WalletRepository Tests +# ============================================================================ + + +class TestWalletRepository: + """Tests for WalletRepository.""" + + @pytest.mark.asyncio + async def test_get_by_address_not_found(self, async_session: AsyncSession) -> None: + """Test getting a non-existent wallet returns None.""" + repo = WalletRepository(async_session) + result = await repo.get_by_address("0xnonexistent") + assert result is None + + @pytest.mark.asyncio + async def test_upsert_creates_new( + self, async_session: AsyncSession, sample_wallet_dto: WalletProfileDTO + ) -> None: + """Test upserting a new wallet profile.""" + repo = WalletRepository(async_session) + await repo.upsert(sample_wallet_dto) + await async_session.commit() + + result = await repo.get_by_address(sample_wallet_dto.address) + assert result is not None + assert result.address == sample_wallet_dto.address.lower() + assert result.nonce == sample_wallet_dto.nonce + assert result.is_fresh == sample_wallet_dto.is_fresh + + @pytest.mark.asyncio + async def test_upsert_updates_existing( + self, async_session: AsyncSession, sample_wallet_dto: WalletProfileDTO + ) -> None: + """Test upserting updates existing profile.""" + repo = WalletRepository(async_session) + await repo.upsert(sample_wallet_dto) + await async_session.commit() + + # Update the DTO + updated_dto = WalletProfileDTO( + address=sample_wallet_dto.address, + nonce=10, + first_seen_at=sample_wallet_dto.first_seen_at, + is_fresh=False, + matic_balance=sample_wallet_dto.matic_balance, + usdc_balance=Decimal("2000.00"), + analyzed_at=datetime.now(UTC), + ) + await repo.upsert(updated_dto) + await async_session.commit() + + result = await repo.get_by_address(sample_wallet_dto.address) + assert result is not None + assert result.nonce == 10 + assert result.is_fresh is False + assert result.usdc_balance == Decimal("2000.00") + + @pytest.mark.asyncio + async def test_get_many( + self, async_session: AsyncSession, sample_wallet_dto: WalletProfileDTO + ) -> None: + """Test getting multiple wallets.""" + repo = WalletRepository(async_session) + + # Insert two wallets + dto2 = WalletProfileDTO( + address="0xbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb", + nonce=3, + first_seen_at=datetime.now(UTC), + is_fresh=True, + matic_balance=None, + usdc_balance=None, + analyzed_at=datetime.now(UTC), + ) + await repo.upsert(sample_wallet_dto) + await repo.upsert(dto2) + await async_session.commit() + + results = await repo.get_many([sample_wallet_dto.address, dto2.address]) + assert len(results) == 2 + + @pytest.mark.asyncio + async def test_get_fresh_wallets( + self, async_session: AsyncSession, sample_wallet_dto: WalletProfileDTO + ) -> None: + """Test getting fresh wallets.""" + repo = WalletRepository(async_session) + + # Insert fresh and non-fresh wallets + non_fresh = WalletProfileDTO( + address="0xcccccccccccccccccccccccccccccccccccccccc", + nonce=100, + first_seen_at=datetime.now(UTC) - timedelta(days=30), + is_fresh=False, + matic_balance=None, + usdc_balance=None, + analyzed_at=datetime.now(UTC), + ) + await repo.upsert(sample_wallet_dto) + await repo.upsert(non_fresh) + await async_session.commit() + + results = await repo.get_fresh_wallets() + assert len(results) == 1 + assert results[0].is_fresh is True + + @pytest.mark.asyncio + async def test_delete( + self, async_session: AsyncSession, sample_wallet_dto: WalletProfileDTO + ) -> None: + """Test deleting a wallet profile.""" + repo = WalletRepository(async_session) + await repo.upsert(sample_wallet_dto) + await async_session.commit() + + deleted = await repo.delete(sample_wallet_dto.address) + await async_session.commit() + assert deleted is True + + result = await repo.get_by_address(sample_wallet_dto.address) + assert result is None + + @pytest.mark.asyncio + async def test_delete_not_found(self, async_session: AsyncSession) -> None: + """Test deleting non-existent wallet returns False.""" + repo = WalletRepository(async_session) + deleted = await repo.delete("0xnonexistent") + assert deleted is False + + @pytest.mark.asyncio + async def test_mark_stale( + self, async_session: AsyncSession, sample_wallet_dto: WalletProfileDTO + ) -> None: + """Test marking a wallet as stale.""" + repo = WalletRepository(async_session) + await repo.upsert(sample_wallet_dto) + await async_session.commit() + + marked = await repo.mark_stale(sample_wallet_dto.address) + await async_session.commit() + assert marked is True + + result = await repo.get_by_address(sample_wallet_dto.address) + assert result is not None + assert result.analyzed_at.year == 2000 + + +# ============================================================================ +# FundingRepository Tests +# ============================================================================ + + +class TestFundingRepository: + """Tests for FundingRepository.""" + + @pytest.mark.asyncio + async def test_insert( + self, async_session: AsyncSession, sample_transfer_dto: FundingTransferDTO + ) -> None: + """Test inserting a funding transfer.""" + repo = FundingRepository(async_session) + await repo.insert(sample_transfer_dto) + await async_session.commit() + + result = await repo.get_by_tx_hash(sample_transfer_dto.tx_hash) + assert result is not None + assert result.amount == sample_transfer_dto.amount + + @pytest.mark.asyncio + async def test_get_transfers_to( + self, async_session: AsyncSession, sample_transfer_dto: FundingTransferDTO + ) -> None: + """Test getting transfers to an address.""" + repo = FundingRepository(async_session) + await repo.insert(sample_transfer_dto) + await async_session.commit() + + results = await repo.get_transfers_to(sample_transfer_dto.to_address) + assert len(results) == 1 + assert results[0].from_address == sample_transfer_dto.from_address.lower() + + @pytest.mark.asyncio + async def test_get_transfers_from( + self, async_session: AsyncSession, sample_transfer_dto: FundingTransferDTO + ) -> None: + """Test getting transfers from an address.""" + repo = FundingRepository(async_session) + await repo.insert(sample_transfer_dto) + await async_session.commit() + + results = await repo.get_transfers_from(sample_transfer_dto.from_address) + assert len(results) == 1 + assert results[0].to_address == sample_transfer_dto.to_address.lower() + + @pytest.mark.asyncio + async def test_get_first_transfer_to( + self, async_session: AsyncSession, sample_transfer_dto: FundingTransferDTO + ) -> None: + """Test getting first transfer to an address.""" + repo = FundingRepository(async_session) + + # Insert multiple transfers with different timestamps + earlier = FundingTransferDTO( + from_address="0xeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee", + to_address=sample_transfer_dto.to_address, + amount=Decimal("100.00"), + token="USDC", + tx_hash="0x" + "b" * 64, + block_number=12345670, + timestamp=datetime.now(UTC) - timedelta(hours=2), + ) + await repo.insert(earlier) + await repo.insert(sample_transfer_dto) + await async_session.commit() + + result = await repo.get_first_transfer_to(sample_transfer_dto.to_address) + assert result is not None + assert result.tx_hash == earlier.tx_hash.lower() + + @pytest.mark.asyncio + async def test_insert_many(self, async_session: AsyncSession) -> None: + """Test inserting multiple transfers.""" + repo = FundingRepository(async_session) + + transfers = [ + FundingTransferDTO( + from_address=f"0x{'a' * 40}", + to_address=f"0x{'b' * 40}", + amount=Decimal(f"{i * 100}.00"), + token="USDC", + tx_hash=f"0x{str(i) * 64}"[:66], + block_number=12345678 + i, + timestamp=datetime.now(UTC), + ) + for i in range(1, 4) + ] + + count = await repo.insert_many(transfers) + await async_session.commit() + assert count == 3 + + +# ============================================================================ +# RelationshipRepository Tests +# ============================================================================ + + +class TestRelationshipRepository: + """Tests for RelationshipRepository.""" + + @pytest.mark.asyncio + async def test_upsert( + self, async_session: AsyncSession, sample_relationship_dto: WalletRelationshipDTO + ) -> None: + """Test upserting a relationship.""" + repo = RelationshipRepository(async_session) + await repo.upsert(sample_relationship_dto) + await async_session.commit() + + results = await repo.get_relationships(sample_relationship_dto.wallet_a) + assert len(results) == 1 + assert results[0].confidence == sample_relationship_dto.confidence + + @pytest.mark.asyncio + async def test_get_relationships_filter_type( + self, async_session: AsyncSession, sample_relationship_dto: WalletRelationshipDTO + ) -> None: + """Test getting relationships with type filter.""" + repo = RelationshipRepository(async_session) + await repo.upsert(sample_relationship_dto) + + same_entity = WalletRelationshipDTO( + wallet_a=sample_relationship_dto.wallet_a, + wallet_b="0xdddddddddddddddddddddddddddddddddddddddd", + relationship_type="same_entity", + confidence=Decimal("0.80"), + ) + await repo.upsert(same_entity) + await async_session.commit() + + funded_results = await repo.get_relationships( + sample_relationship_dto.wallet_a, relationship_type="funded_by" + ) + assert len(funded_results) == 1 + assert funded_results[0].relationship_type == "funded_by" + + @pytest.mark.asyncio + async def test_get_related_wallets( + self, async_session: AsyncSession, sample_relationship_dto: WalletRelationshipDTO + ) -> None: + """Test getting related wallet addresses.""" + repo = RelationshipRepository(async_session) + await repo.upsert(sample_relationship_dto) + await async_session.commit() + + related = await repo.get_related_wallets(sample_relationship_dto.wallet_a) + assert len(related) == 1 + assert sample_relationship_dto.wallet_b.lower() in related + + @pytest.mark.asyncio + async def test_delete_relationship( + self, async_session: AsyncSession, sample_relationship_dto: WalletRelationshipDTO + ) -> None: + """Test deleting a relationship.""" + repo = RelationshipRepository(async_session) + await repo.upsert(sample_relationship_dto) + await async_session.commit() + + deleted = await repo.delete( + sample_relationship_dto.wallet_a, + sample_relationship_dto.wallet_b, + sample_relationship_dto.relationship_type, + ) + await async_session.commit() + assert deleted is True + + results = await repo.get_relationships(sample_relationship_dto.wallet_a) + assert len(results) == 0 + + @pytest.mark.asyncio + async def test_upsert_updates_confidence( + self, async_session: AsyncSession, sample_relationship_dto: WalletRelationshipDTO + ) -> None: + """Test that upserting updates the confidence.""" + repo = RelationshipRepository(async_session) + await repo.upsert(sample_relationship_dto) + await async_session.commit() + + updated = WalletRelationshipDTO( + wallet_a=sample_relationship_dto.wallet_a, + wallet_b=sample_relationship_dto.wallet_b, + relationship_type=sample_relationship_dto.relationship_type, + confidence=Decimal("0.99"), + ) + await repo.upsert(updated) + await async_session.commit() + + results = await repo.get_relationships(sample_relationship_dto.wallet_a) + assert len(results) == 1 + assert results[0].confidence == Decimal("0.99") + + +# ============================================================================ +# DTO Tests +# ============================================================================ + + +class TestDTOs: + """Tests for Data Transfer Objects.""" + + def test_wallet_profile_dto_from_model(self) -> None: + """Test WalletProfileDTO.from_model works correctly.""" + from polymarket_insider_tracker.storage.models import WalletProfileModel + + now = datetime.now(UTC) + model = WalletProfileModel( + id=1, + address="0x1234", + nonce=5, + first_seen_at=now, + is_fresh=True, + matic_balance=Decimal("100"), + usdc_balance=Decimal("50.00"), + analyzed_at=now, + created_at=now, + updated_at=now, + ) + + dto = WalletProfileDTO.from_model(model) + assert dto.address == "0x1234" + assert dto.nonce == 5 + assert dto.is_fresh is True + + def test_funding_transfer_dto_from_model(self) -> None: + """Test FundingTransferDTO.from_model works correctly.""" + from polymarket_insider_tracker.storage.models import FundingTransferModel + + now = datetime.now(UTC) + model = FundingTransferModel( + id=1, + from_address="0xaaa", + to_address="0xbbb", + amount=Decimal("100.00"), + token="USDC", + tx_hash="0x123", + block_number=12345, + timestamp=now, + created_at=now, + ) + + dto = FundingTransferDTO.from_model(model) + assert dto.from_address == "0xaaa" + assert dto.amount == Decimal("100.00") + + def test_wallet_relationship_dto_from_model(self) -> None: + """Test WalletRelationshipDTO.from_model works correctly.""" + from polymarket_insider_tracker.storage.models import WalletRelationshipModel + + now = datetime.now(UTC) + model = WalletRelationshipModel( + id=1, + wallet_a="0xaaa", + wallet_b="0xbbb", + relationship_type="funded_by", + confidence=Decimal("0.95"), + created_at=now, + ) + + dto = WalletRelationshipDTO.from_model(model) + assert dto.wallet_a == "0xaaa" + assert dto.relationship_type == "funded_by" + assert dto.confidence == Decimal("0.95")