test_basic.py•7.78 kB
#!/usr/bin/env python3
"""
Basic test script for plotnine MCP server functionality.
"""
import sys
import pandas as pd
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.plot_builder import build_plot, save_plot
from plotnine_mcp.schemas import (
Aesthetics,
GeomConfig,
ThemeConfig,
LabelsConfig,
OutputConfig,
)
def test_inline_data_scatter_plot():
"""Test creating a simple scatter plot with inline data."""
print("Test 1: Simple scatter plot with inline data...")
# Create data source
data_source = DataSource(
type="inline",
data=[
{"x": 1, "y": 2, "category": "A"},
{"x": 2, "y": 4, "category": "A"},
{"x": 3, "y": 3, "category": "B"},
{"x": 4, "y": 5, "category": "B"},
{"x": 5, "y": 6, "category": "C"},
],
)
# Load data
data = load_data(data_source)
print(f" ✓ Data loaded: {len(data)} rows")
# Create plot configuration
aes_config = Aesthetics(x="x", y="y", color="category")
geom_config = GeomConfig(type="point", params={"size": 3})
theme_config = ThemeConfig(base="minimal")
labels_config = LabelsConfig(
title="Test Scatter Plot", x="X Values", y="Y Values"
)
output_config = OutputConfig(
format="png", filename="test_scatter.png", directory="./output"
)
# Build plot
plot = build_plot(
data=data,
aes_config=aes_config,
geom_config=geom_config,
theme_config=theme_config,
labels_config=labels_config,
)
print(" ✓ Plot built successfully")
# Save plot
result = save_plot(plot, output_config)
print(f" ✓ Plot saved to: {result['path']}")
return result
def test_file_data_line_plot():
"""Test creating a line plot from CSV file."""
print("\nTest 2: Line plot from CSV file...")
# Create data source
data_source = DataSource(
type="file", path="./examples/sample_data.csv", format="csv"
)
# Load data
data = load_data(data_source)
print(f" ✓ Data loaded: {len(data)} rows, columns: {list(data.columns)}")
# Create plot configuration
aes_config = Aesthetics(x="x", y="y", color="category")
geom_config = GeomConfig(type="line", params={"size": 1.5})
theme_config = ThemeConfig(
base="bw",
customizations={"figure_size": [10, 6], "legend_position": "bottom"},
)
labels_config = LabelsConfig(
title="Sample Data Line Plot", x="X Axis", y="Y Axis"
)
output_config = OutputConfig(
format="png", filename="test_line.png", directory="./output"
)
# Build plot
plot = build_plot(
data=data,
aes_config=aes_config,
geom_config=geom_config,
theme_config=theme_config,
labels_config=labels_config,
)
print(" ✓ Plot built successfully")
# Save plot
result = save_plot(plot, output_config)
print(f" ✓ Plot saved to: {result['path']}")
return result
def test_bar_plot():
"""Test creating a bar plot."""
print("\nTest 3: Bar plot...")
# Create data
data_source = DataSource(
type="inline",
data=[
{"category": "A", "value": 10},
{"category": "B", "value": 15},
{"category": "C", "value": 12},
{"category": "D", "value": 18},
],
)
data = load_data(data_source)
print(f" ✓ Data loaded: {len(data)} rows")
# Create plot
aes_config = Aesthetics(x="category", y="value", fill="category")
geom_config = GeomConfig(type="col")
theme_config = ThemeConfig(base="classic")
labels_config = LabelsConfig(title="Bar Chart Example", x="Category", y="Value")
output_config = OutputConfig(
format="png", filename="test_bar.png", directory="./output"
)
plot = build_plot(
data=data,
aes_config=aes_config,
geom_config=geom_config,
theme_config=theme_config,
labels_config=labels_config,
)
print(" ✓ Plot built successfully")
result = save_plot(plot, output_config)
print(f" ✓ Plot saved to: {result['path']}")
return result
def test_multi_layer_plot():
"""Test creating a multi-layer plot (scatter + smooth)."""
print("\nTest 4: Multi-layer plot (scatter + smooth)...")
# Create data source
data_source = DataSource(
type="file", path="./examples/sample_data.csv", format="csv"
)
data = load_data(data_source)
print(f" ✓ Data loaded: {len(data)} rows")
# Create plot with multiple geoms
aes_config = Aesthetics(x="x", y="y", color="category")
geom_configs = [
GeomConfig(type="point", params={"size": 3, "alpha": 0.7}),
GeomConfig(type="smooth", params={"method": "lm", "se": False}),
]
theme_config = ThemeConfig(
base="minimal", customizations={"figure_size": [10, 6]}
)
labels_config = LabelsConfig(
title="Multi-Layer Plot: Points + Smooth Trend",
x="X Values",
y="Y Values",
)
output_config = OutputConfig(
format="png", filename="test_multi_layer.png", directory="./output"
)
plot = build_plot(
data=data,
aes_config=aes_config,
geom_configs=geom_configs,
theme_config=theme_config,
labels_config=labels_config,
)
print(" ✓ Multi-layer plot built successfully")
result = save_plot(plot, output_config)
print(f" ✓ Plot saved to: {result['path']}")
return result
def test_boxplot_with_jitter():
"""Test boxplot with jittered points overlay."""
print("\nTest 5: Boxplot with jittered points...")
# Create data
data_source = DataSource(
type="file", path="./examples/sample_data.csv", format="csv"
)
data = load_data(data_source)
print(f" ✓ Data loaded: {len(data)} rows")
# Create layered plot
aes_config = Aesthetics(x="category", y="y", fill="category")
geom_configs = [
GeomConfig(type="boxplot", params={"alpha": 0.7}),
GeomConfig(type="jitter", params={"width": 0.2, "alpha": 0.5}),
]
theme_config = ThemeConfig(base="bw")
labels_config = LabelsConfig(
title="Boxplot with Individual Points",
x="Category",
y="Values",
)
output_config = OutputConfig(
format="png", filename="test_boxplot_jitter.png", directory="./output"
)
plot = build_plot(
data=data,
aes_config=aes_config,
geom_configs=geom_configs,
theme_config=theme_config,
labels_config=labels_config,
)
print(" ✓ Boxplot with jitter built successfully")
result = save_plot(plot, output_config)
print(f" ✓ Plot saved to: {result['path']}")
return result
def main():
"""Run all tests."""
print("=" * 60)
print("Running Plotnine MCP Basic Tests")
print("=" * 60)
try:
# Test 1: Scatter plot with inline data
test_inline_data_scatter_plot()
# Test 2: Line plot from file
test_file_data_line_plot()
# Test 3: Bar plot
test_bar_plot()
# Test 4: Multi-layer plot
test_multi_layer_plot()
# Test 5: Boxplot with jitter
test_boxplot_with_jitter()
print("\n" + "=" * 60)
print("All tests passed! ✓")
print("=" * 60)
print("\nCheck the ./output directory for generated plots.")
except Exception as e:
print(f"\n✗ Test failed with error: {e}")
import traceback
traceback.print_exc()
sys.exit(1)
if __name__ == "__main__":
main()