MaverickMCP

""" Mock data persistence implementation for testing. """ from datetime import datetime from typing import Any import pandas as pd from sqlalchemy.orm import Session class MockSession: """Mock SQLAlchemy session for testing.""" def __init__(self): self.closed = False self.committed = False self.rolled_back = False def close(self): self.closed = True def commit(self): self.committed = True def rollback(self): self.rolled_back = True class MockDataPersistence: """ Mock implementation of IDataPersistence for testing. """ def __init__(self): """Initialize the mock persistence layer.""" self._price_data: dict[str, pd.DataFrame] = {} self._stock_records: dict[str, dict[str, Any]] = {} self._screening_results: dict[str, list[dict[str, Any]]] = {} self._call_log: list[dict[str, Any]] = [] async def get_session(self) -> Session: """Get a mock database session.""" self._log_call("get_session", {}) return MockSession() async def get_read_only_session(self) -> Session: """Get a mock read-only database session.""" self._log_call("get_read_only_session", {}) return MockSession() async def save_price_data( self, session: Session, symbol: str, data: pd.DataFrame ) -> int: """Save mock price data.""" self._log_call("save_price_data", {"symbol": symbol, "data_length": len(data)}) symbol = symbol.upper() # Store the data if symbol in self._price_data: # Merge with existing data existing_data = self._price_data[symbol] combined = pd.concat([existing_data, data]) # Remove duplicates and sort combined = combined[~combined.index.duplicated(keep="last")].sort_index() self._price_data[symbol] = combined else: self._price_data[symbol] = data.copy() return len(data) async def get_price_data( self, session: Session, symbol: str, start_date: str, end_date: str, ) -> pd.DataFrame: """Retrieve mock price data.""" self._log_call( "get_price_data", { "symbol": symbol, "start_date": start_date, "end_date": end_date, }, ) symbol = symbol.upper() if symbol not in self._price_data: # Return empty DataFrame with proper columns return pd.DataFrame(columns=["Open", "High", "Low", "Close", "Volume"]) data = self._price_data[symbol].copy() # Filter by date range if start_date: data = data[data.index >= start_date] if end_date: data = data[data.index <= end_date] return data async def get_or_create_stock(self, session: Session, symbol: str) -> Any: """Get or create a mock stock record.""" self._log_call("get_or_create_stock", {"symbol": symbol}) symbol = symbol.upper() if symbol not in self._stock_records: self._stock_records[symbol] = { "symbol": symbol, "company_name": f"{symbol} Inc.", "sector": "Technology", "industry": "Software", "exchange": "NASDAQ", "currency": "USD", "country": "US", "created_at": datetime.now(), } return self._stock_records[symbol] async def save_screening_results( self, session: Session, screening_type: str, results: list[dict[str, Any]], ) -> int: """Save mock screening results.""" self._log_call( "save_screening_results", { "screening_type": screening_type, "results_count": len(results), }, ) self._screening_results[screening_type] = results.copy() return len(results) async def get_screening_results( self, session: Session, screening_type: str, limit: int | None = None, min_score: float | None = None, ) -> list[dict[str, Any]]: """Retrieve mock screening results.""" self._log_call( "get_screening_results", { "screening_type": screening_type, "limit": limit, "min_score": min_score, }, ) if screening_type not in self._screening_results: return self._generate_default_screening_results(screening_type) results = self._screening_results[screening_type].copy() # Apply filters if min_score is not None: if screening_type == "maverick": results = [ r for r in results if r.get("combined_score", 0) >= min_score ] elif screening_type == "bearish": results = [r for r in results if r.get("score", 0) >= min_score] elif screening_type == "trending": results = [ r for r in results if r.get("momentum_score", 0) >= min_score ] if limit is not None: results = results[:limit] return results async def get_latest_screening_data(self) -> dict[str, list[dict[str, Any]]]: """Get mock latest screening data.""" self._log_call("get_latest_screening_data", {}) return { "maverick_stocks": await self.get_screening_results(None, "maverick"), "maverick_bear_stocks": await self.get_screening_results(None, "bearish"), "supply_demand_breakouts": await self.get_screening_results( None, "trending" ), } async def check_data_freshness(self, symbol: str, max_age_hours: int = 24) -> bool: """Check mock data freshness.""" self._log_call( "check_data_freshness", {"symbol": symbol, "max_age_hours": max_age_hours} ) # For testing, assume data is fresh if it exists symbol = symbol.upper() return symbol in self._price_data async def bulk_save_price_data( self, session: Session, symbol: str, data: pd.DataFrame ) -> int: """Bulk save mock price data.""" return await self.save_price_data(session, symbol, data) async def get_symbols_with_data( self, session: Session, limit: int | None = None ) -> list[str]: """Get mock list of symbols with data.""" self._log_call("get_symbols_with_data", {"limit": limit}) symbols = list(self._price_data.keys()) if limit is not None: symbols = symbols[:limit] return symbols async def cleanup_old_data(self, session: Session, days_to_keep: int = 365) -> int: """Mock cleanup of old data.""" self._log_call("cleanup_old_data", {"days_to_keep": days_to_keep}) # For testing, return 0 (no cleanup performed) return 0 def _generate_default_screening_results( self, screening_type: str ) -> list[dict[str, Any]]: """Generate default screening results for testing.""" if screening_type == "maverick": return [ { "symbol": "TEST1", "combined_score": 95, "momentum_score": 92, "pattern": "Cup with Handle", "consolidation": "yes", "squeeze": "firing", }, { "symbol": "TEST2", "combined_score": 88, "momentum_score": 85, "pattern": "Flat Base", "consolidation": "no", "squeeze": "setup", }, ] elif screening_type == "bearish": return [ { "symbol": "BEAR1", "score": 92, "momentum_score": 25, "rsi_14": 28, "atr_contraction": True, "big_down_vol": True, }, ] elif screening_type == "trending": return [ { "symbol": "TREND1", "momentum_score": 95, "close": 150.25, "sma_50": 145.50, "sma_150": 140.25, "sma_200": 135.75, "pattern": "Breakout", }, ] else: return [] # Testing utilities def _log_call(self, method: str, args: dict[str, Any]) -> None: """Log method calls for testing verification.""" self._call_log.append( { "method": method, "args": args, "timestamp": datetime.now(), } ) def get_call_log(self) -> list[dict[str, Any]]: """Get the log of method calls.""" return self._call_log.copy() def clear_call_log(self) -> None: """Clear the method call log.""" self._call_log.clear() def set_price_data(self, symbol: str, data: pd.DataFrame) -> None: """Set price data for testing.""" self._price_data[symbol.upper()] = data def get_stored_price_data(self, symbol: str) -> pd.DataFrame | None: """Get stored price data for testing verification.""" return self._price_data.get(symbol.upper()) def set_screening_results( self, screening_type: str, results: list[dict[str, Any]] ) -> None: """Set screening results for testing.""" self._screening_results[screening_type] = results def clear_all_data(self) -> None: """Clear all stored data.""" self._price_data.clear() self._stock_records.clear() self._screening_results.clear()