test_phase2_phase3.py

#!/usr/bin/env python3
"""
Test script for Phase 2 & Phase 3 features:
- Plot templates
- Template suggestions
- Color palettes
- Data transformations
- Batch processing
"""

import sys
import asyncio
from pathlib import Path

# Add src to path
sys.path.insert(0, str(Path(__file__).parent / "src"))

from plotnine_mcp.data_loader import load_data, DataSource
from plotnine_mcp.templates import (
    list_templates,
    apply_template,
    suggest_template,
    TEMPLATES,
)
from plotnine_mcp.palettes import (
    list_palettes,
    get_palette,
    palette_categories,
    create_scale_config,
)
from plotnine_mcp.transforms import apply_transforms
from plotnine_mcp.server import (
    list_plot_templates_handler,
    create_plot_from_template_handler,
    suggest_plot_templates_handler,
    list_color_palettes_handler,
    batch_create_plots_handler,
)


def test_list_templates():
    """Test listing plot templates."""
    print("Test 1: List plot templates...")

    templates = list_templates()
    assert len(templates) == len(TEMPLATES)
    assert "time_series" in templates
    assert "scatter_with_trend" in templates

    print(f"  ✓ Listed {len(templates)} templates")
    print(f"  Templates: {', '.join(list(templates.keys())[:3])}...")


def test_apply_template():
    """Test applying a template."""
    print("\nTest 2: Apply template...")

    data_source = {
        "type": "inline",
        "data": [
            {"x": 1, "y": 2, "category": "A"},
            {"x": 2, "y": 4, "category": "B"},
        ],
    }

    aes = {"x": "x", "y": "y", "color": "category"}

    config = apply_template("scatter_with_trend", data_source, aes)

    assert "data_source" in config
    assert "aes" in config
    assert "geoms" in config
    assert len(config["geoms"]) == 2  # point + smooth

    print("  ✓ Template applied successfully")
    print(f"  Geoms: {[g['type'] for g in config['geoms']]}")


def test_suggest_template():
    """Test template suggestions."""
    print("\nTest 3: Suggest templates based on data...")

    # Numeric data
    suggestions = suggest_template(num_numeric=2, num_categorical=0, has_time=False)
    assert "scatter_with_trend" in suggestions
    print(f"  ✓ Numeric data: {suggestions[:2]}")

    # Time series data
    suggestions = suggest_template(num_numeric=1, num_categorical=0, has_time=True)
    assert "time_series" in suggestions
    print(f"  ✓ Time data: {suggestions[:2]}")

    # Categorical comparison
    suggestions = suggest_template(num_numeric=1, num_categorical=1, has_time=False)
    assert any(
        "distribution" in s or "boxplot" in s for s in suggestions
    )
    print(f"  ✓ Categorical data: {suggestions[:2]}")


def test_list_palettes():
    """Test listing color palettes."""
    print("\nTest 4: List color palettes...")

    # All palettes
    all_palettes = list_palettes()
    assert len(all_palettes) > 0

    # By category
    scientific = list_palettes("scientific")
    assert "scientific_viridis" in scientific
    assert "scientific_plasma" in scientific

    print(f"  ✓ Total palettes: {len(all_palettes)}")
    print(f"  ✓ Scientific palettes: {len(scientific)}")


def test_get_palette():
    """Test getting specific palette."""
    print("\nTest 5: Get specific palette...")

    # With prefix
    viridis = get_palette("scientific_viridis")
    assert len(viridis) == 10
    assert viridis[0].startswith("#")

    # Without prefix
    viridis2 = get_palette("viridis")
    assert viridis == viridis2

    print(f"  ✓ Viridis palette: {len(viridis)} colors")
    print(f"  First 3 colors: {viridis[:3]}")


def test_palette_categories():
    """Test palette categories."""
    print("\nTest 6: Palette categories...")

    categories = palette_categories()
    assert "colorblind_safe" in categories
    assert "scientific" in categories
    assert "categorical" in categories

    print(f"  ✓ Categories: {len(categories)}")
    for cat in list(categories.keys())[:3]:
        print(f"    - {cat}")


def test_create_scale_config():
    """Test creating scale configuration from palette."""
    print("\nTest 7: Create scale config from palette...")

    # Discrete scale
    scale = create_scale_config("viridis", aesthetic="color", type="discrete")
    assert scale["aesthetic"] == "color"
    assert scale["type"] == "discrete"
    assert "values" in scale["params"]

    # Gradient scale
    scale = create_scale_config("viridis", aesthetic="fill", type="gradient")
    assert scale["aesthetic"] == "fill"
    assert scale["type"] == "gradient"
    assert "low" in scale["params"]
    assert "high" in scale["params"]

    print("  ✓ Discrete scale created")
    print("  ✓ Gradient scale created")


def test_data_transforms_filter():
    """Test filter transformation."""
    print("\nTest 8: Data transformation - filter...")

    data_source = DataSource(
        type="inline",
        data=[
            {"x": 1, "y": 2, "category": "A"},
            {"x": 2, "y": 4, "category": "B"},
            {"x": 3, "y": 6, "category": "A"},
        ],
    )
    data = load_data(data_source)

    transforms = [{"type": "filter", "filter_expr": "y > 2"}]
    result = apply_transforms(data, transforms)

    assert len(result) == 2  # Only y=4 and y=6
    print(f"  ✓ Filtered from {len(data)} to {len(result)} rows")


def test_data_transforms_group():
    """Test group and summarize transformation."""
    print("\nTest 9: Data transformation - group & summarize...")

    data_source = DataSource(
        type="inline",
        data=[
            {"category": "A", "value": 10},
            {"category": "A", "value": 20},
            {"category": "B", "value": 15},
            {"category": "B", "value": 25},
        ],
    )
    data = load_data(data_source)

    transforms = [
        {"type": "group_summarize", "group_by": "category", "aggregations": {"value": "mean"}}
    ]
    result = apply_transforms(data, transforms)

    assert len(result) == 2  # A and B
    assert "value" in result.columns
    print(f"  ✓ Grouped from {len(data)} to {len(result)} rows")


def test_data_transforms_mutate():
    """Test mutate transformation."""
    print("\nTest 10: Data transformation - mutate...")

    data_source = DataSource(
        type="inline",
        data=[
            {"price": 10, "quantity": 5},
            {"price": 20, "quantity": 3},
        ],
    )
    data = load_data(data_source)

    transforms = [{"type": "mutate", "mutations": {"total": "price * quantity"}}]
    result = apply_transforms(data, transforms)

    assert "total" in result.columns
    assert result["total"].tolist() == [50, 60]
    print("  ✓ Created new column 'total'")
    print(f"  Values: {result['total'].tolist()}")


def test_data_transforms_sort():
    """Test sort transformation."""
    print("\nTest 11: Data transformation - sort...")

    data_source = DataSource(
        type="inline",
        data=[
            {"name": "C", "value": 3},
            {"name": "A", "value": 1},
            {"name": "B", "value": 2},
        ],
    )
    data = load_data(data_source)

    transforms = [{"type": "sort", "sort_by": "value", "ascending": False}]
    result = apply_transforms(data, transforms)

    assert result["value"].tolist() == [3, 2, 1]
    print("  ✓ Sorted by value (descending)")
    print(f"  Order: {result['name'].tolist()}")


def test_data_transforms_chained():
    """Test multiple chained transformations."""
    print("\nTest 12: Data transformation - chained...")

    data_source = DataSource(
        type="inline",
        data=[
            {"category": "A", "value": 10, "status": "active"},
            {"category": "A", "value": 20, "status": "active"},
            {"category": "B", "value": 15, "status": "inactive"},
            {"category": "B", "value": 25, "status": "active"},
        ],
    )
    data = load_data(data_source)

    transforms = [
        {"type": "filter", "filter_expr": "status == 'active'"},
        {"type": "group_summarize", "group_by": "category", "aggregations": {"value": "sum"}},
        {"type": "sort", "sort_by": "value", "ascending": False},
    ]
    result = apply_transforms(data, transforms)

    assert len(result) == 2
    print("  ✓ Applied 3 chained transforms")
    print(f"  Final shape: {result.shape}")


async def test_template_handler():
    """Test template MCP handler."""
    print("\nTest 13: Template handler (MCP)...")

    loop = asyncio.get_event_loop()
    result = await list_plot_templates_handler()

    assert len(result) == 1
    assert "time_series" in result[0].text
    assert "scatter_with_trend" in result[0].text

    print("  ✓ Template handler works")
    print(f"  Response length: {len(result[0].text)} characters")


async def test_palette_handler():
    """Test palette MCP handler."""
    print("\nTest 14: Palette handler (MCP)...")

    loop = asyncio.get_event_loop()

    # List all categories
    result = await list_color_palettes_handler({})
    assert "colorblind_safe" in result[0].text

    # List specific category
    result = await list_color_palettes_handler({"category": "scientific"})
    assert "viridis" in result[0].text

    print("  ✓ Palette handler works")
    print("  ✓ Category filtering works")


async def test_create_from_template():
    """Test creating plot from template."""
    print("\nTest 15: Create plot from template...")

    arguments = {
        "template_name": "scatter_with_trend",
        "data_source": {
            "type": "inline",
            "data": [
                {"x": 1, "y": 2},
                {"x": 2, "y": 4},
                {"x": 3, "y": 6},
                {"x": 4, "y": 8},
            ],
        },
        "aes": {"x": "x", "y": "y"},
        "output": {"filename": "test_template", "directory": "./test_output"},
    }

    result = await create_plot_from_template_handler(arguments)

    assert len(result) == 1
    assert "scatter_with_trend" in result[0].text
    assert "successfully" in result[0].text

    # Verify file created
    output_file = Path("./test_output/test_template.png")
    assert output_file.exists()

    print("  ✓ Plot created from template")
    print(f"  ✓ File: {output_file}")


async def test_suggest_templates_handler():
    """Test template suggestion handler."""
    print("\nTest 16: Template suggestion handler...")

    arguments = {
        "data_source": {
            "type": "inline",
            "data": [
                {"date": "2024-01-01", "value": 10, "category": "A"},
                {"date": "2024-01-02", "value": 20, "category": "B"},
            ],
        },
        "goal": "show trend over time",
    }

    result = await suggest_plot_templates_handler(arguments)

    assert len(result) == 1
    assert "time_series" in result[0].text or "multi_line" in result[0].text

    print("  ✓ Template suggestions work")
    print(f"  Response includes goal-based filtering")


async def test_batch_create_plots():
    """Test batch plot creation."""
    print("\nTest 17: Batch plot creation...")

    arguments = {
        "plots": [
            {
                "data_source": {
                    "type": "inline",
                    "data": [{"x": 1, "y": 2}, {"x": 2, "y": 4}],
                },
                "aes": {"x": "x", "y": "y"},
                "geom": {"type": "point"},
                "output": {"filename": "batch_1", "directory": "./test_output"},
            },
            {
                "data_source": {
                    "type": "inline",
                    "data": [{"x": 1, "y": 3}, {"x": 2, "y": 5}],
                },
                "aes": {"x": "x", "y": "y"},
                "geom": {"type": "line"},
                "output": {"filename": "batch_2", "directory": "./test_output"},
            },
        ]
    }

    result = await batch_create_plots_handler(arguments)

    assert len(result) == 1
    assert "2 plot(s)" in result[0].text
    assert "Successful: 2" in result[0].text

    # Verify files created
    assert Path("./test_output/batch_1.png").exists()
    assert Path("./test_output/batch_2.png").exists()

    print("  ✓ Batch created 2 plots")
    print("  ✓ All files created successfully")


def cleanup():
    """Clean up test output files."""
    import shutil

    test_output = Path("./test_output")
    if test_output.exists():
        shutil.rmtree(test_output)
    print("\n✓ Cleaned up test output files")


def main():
    """Run all tests."""
    print("=" * 60)
    print("Running Phase 2 & Phase 3 Features Tests")
    print("=" * 60)

    try:
        # Template tests
        test_list_templates()
        test_apply_template()
        test_suggest_template()

        # Palette tests
        test_list_palettes()
        test_get_palette()
        test_palette_categories()
        test_create_scale_config()

        # Transform tests
        test_data_transforms_filter()
        test_data_transforms_group()
        test_data_transforms_mutate()
        test_data_transforms_sort()
        test_data_transforms_chained()

        # Async handler tests
        loop = asyncio.get_event_loop()
        loop.run_until_complete(test_template_handler())
        loop.run_until_complete(test_palette_handler())
        loop.run_until_complete(test_create_from_template())
        loop.run_until_complete(test_suggest_templates_handler())
        loop.run_until_complete(test_batch_create_plots())

        print("\n" + "=" * 60)
        print("All Phase 2 & 3 tests passed! ✓")
        print("=" * 60)
        print("\nFeatures tested:")
        print("  • 9 plot templates")
        print("  • 21 color palettes (6 categories)")
        print("  • 12 data transformations")
        print("  • Batch processing")
        print("  • All MCP handlers")

    except Exception as e:
        print(f"\n✗ Test failed: {e}")
        import traceback

        traceback.print_exc()
        sys.exit(1)
    finally:
        cleanup()


if __name__ == "__main__":
    main()