Obsidian Elite RAG MCP Server

"""Tests for dataset first look prompt functionality.""" import pytest import pandas as pd import tempfile import os from mcp_server.prompts.dataset_first_look_prompt import dataset_first_look from mcp_server.models.schemas import DatasetManager, loaded_datasets, dataset_schemas @pytest.fixture def sample_dataset(): """Create a sample dataset for testing.""" data = { 'order_id': ['ord_001', 'ord_002', 'ord_003', 'ord_004', 'ord_005'], 'customer_id': ['cust_123', 'cust_124', 'cust_125', 'cust_126', 'cust_127'], 'product_category': ['electronics', 'books', 'clothing', 'electronics', 'home_garden'], 'order_value': [299.99, 29.99, 89.50, 599.99, 149.99], 'order_date': ['2024-11-15', '2024-11-14', '2024-11-13', '2024-11-12', '2024-11-11'], 'region': ['west_coast', 'midwest', 'east_coast', 'west_coast', 'south'], 'customer_segment': ['premium', 'standard', 'premium', 'premium', 'standard'] } df = pd.DataFrame(data) with tempfile.NamedTemporaryFile(mode='w', suffix='.csv', delete=False) as f: df.to_csv(f.name, index=False) yield f.name # Cleanup os.unlink(f.name) @pytest.fixture def simple_dataset(): """Create a simple dataset with limited columns for testing edge cases.""" data = { 'id': [1, 2, 3, 4, 5], 'name': ['Alice', 'Bob', 'Charlie', 'Diana', 'Eve'] } df = pd.DataFrame(data) with tempfile.NamedTemporaryFile(mode='w', suffix='.csv', delete=False) as f: df.to_csv(f.name, index=False) yield f.name # Cleanup os.unlink(f.name) @pytest.fixture(autouse=True) def clear_datasets(): """Clear datasets before and after each test.""" loaded_datasets.clear() dataset_schemas.clear() yield loaded_datasets.clear() dataset_schemas.clear() class TestDatasetFirstLook: """Test dataset first look prompt functionality.""" @pytest.mark.asyncio async def test_dataset_first_look_comprehensive(self, sample_dataset): """Test first look prompt with comprehensive dataset.""" DatasetManager.load_dataset(sample_dataset, 'ecommerce') result = await dataset_first_look('ecommerce') assert isinstance(result, str) assert 'ecommerce' in result assert '5 records' in result assert 'columns' in result assert '📊 Numerical columns' in result assert '🏷️ Categorical columns' in result # Date columns might be detected as identifiers if all dates are unique assert ('📅 Date/Time columns' in result or '🔑 Identifier columns' in result) assert '🎯 Recommended starting points' in result assert 'find_correlations' in result or 'segment_by_column' in result @pytest.mark.asyncio async def test_dataset_first_look_simple(self, simple_dataset): """Test first look prompt with simple dataset.""" DatasetManager.load_dataset(simple_dataset, 'simple') result = await dataset_first_look('simple') assert isinstance(result, str) assert 'simple' in result assert '5 records' in result # Should still provide useful guidance even with limited data assert 'What aspect' in result @pytest.mark.asyncio async def test_dataset_first_look_nonexistent(self): """Test error handling for non-existent dataset.""" result = await dataset_first_look('nonexistent') assert isinstance(result, str) assert 'not loaded' in result assert 'load_dataset()' in result if __name__ == '__main__': pytest.main([__file__])