Added unit tests

Author: aponcedeleonch

CLAUDE.md

@@ -10,4 +10,9 @@
 - After significant code changes run `make format` and `make typecheck` to make sure the code follows best practices
 
 ## Code Maintenance Guidelines
-- When there is a new plot type added in plot.py we need to update the mcp tool in server.py
+- When there is a new plot type added in plot.py we need to update the mcp tool in server.py
+
+## Dependency Management
+- This project uses uv for managing dependencies. In case you need to add a dependency use the command `uv add <dependency>`
+- If the dependency is for developing, e.g. for testing use `uv add <dependency> --dev`
+- All of the configuration of the project should be centralized in the pyproject.toml file

Makefile

@@ -1,22 +1,25 @@
-.PHONY: install format lint typecheck help
+.PHONY: install format lint typecheck test help
 
 help:
 	@echo "Available commands:"
 	@echo "  install    - Install dependencies using uv"
 	@echo "  format     - Format code using ruff"
 	@echo "  lint       - Run linting using ruff"
 	@echo "  typecheck  - Run type checking using ty"
+	@echo "  test       - Run tests using pytest"
 	@echo "  help       - Show this help message"
 
 install:
 	uv sync
 
 format:
 	uv run ruff format .
-	uv run ruff check --fix .
 
 lint:
 	uv run ruff check .
 
 typecheck:
 	uv run ty check
+
+test:
+	uv run pytest

pyproject.toml

@@ -14,7 +14,7 @@ requires-python = ">=3.13"
 classifiers = [
     "Development Status :: 3 - Alpha",
     "Intended Audience :: Developers",
-    "License :: OSI Approved :: MIT License",
+    "License :: OSI Approved :: Apache Software License",
     "Programming Language :: Python :: 3",
     "Programming Language :: Python :: 3.13",
 ]
@@ -32,6 +32,7 @@ dependencies = [
 
 [dependency-groups]
 dev = [
+    "pytest>=8.4.1",
     "ruff>=0.12.5",
     "ty>=0.0.1a16",
 ]
@@ -52,5 +53,17 @@ lint.select = [
     "I", # isort
     "C", # flake8-comprehensions
     "B", # flake8-bugbear
-    "S", # flake8-bandit
 ]
+lint.ignore = []
+
+[tool.pytest.ini_options]
+minversion = "8.0"
+addopts = [
+    "-ra",
+    "--strict-config",
+    "-v"
+]
+testpaths = ["tests"]
+python_files = ["test_*.py"]
+python_classes = ["Test*"]
+python_functions = ["test_*"]

src/plotting_mcp/plot.py

@@ -45,7 +45,7 @@ def _auto_rotate_labels(ax: plt.Axes, axis: Literal["x", "y"] = "x") -> None:
         ax.tick_params(axis=axis, labelrotation=90)
 
 
-def _create_world_map(ax: GeoAxes, df: pd.DataFrame, **kwargs) -> None:
+def _create_world_map_plot(ax: GeoAxes, df: pd.DataFrame, **kwargs) -> None:
     """Create a world map with coordinate points."""
     # Add map features
     ax.add_feature(cfeature.COASTLINE)
@@ -135,13 +135,17 @@ def _create_pie_plot(ax: plt.Axes, df: pd.DataFrame, **kwargs) -> None:
         )
 
 
-def _create_matplotlib_plot(  # noqa: C901
+def _create_plot(  # noqa: C901
     df: pd.DataFrame, plot_type: str, **kwargs
 ) -> tuple[plt.Figure, plt.Axes]:
     """Create a plot using matplotlib/seaborn."""
     if df.empty:
         raise ValueError("CSV data is empty")
 
+    # Validate that the DataFrame contains no NaN values
+    if df.isnull().any().any():
+        raise ValueError("CSV data contains NaN/null values. Please ensure all data is complete.")
+
     supported_plot_types = ["line", "bar", "pie", "worldmap"]
     if plot_type not in supported_plot_types:
         raise ValueError(
@@ -169,7 +173,7 @@ def _create_matplotlib_plot(  # noqa: C901
         _create_pie_plot(ax, df, **kwargs)
     elif plot_type == "worldmap":
         # Cartopy doesn't return correct Axes type, so we ignore type checking
-        _create_world_map(ax, df, **kwargs)  # ty: ignore[invalid-argument-type]
+        _create_world_map_plot(ax, df, **kwargs)  # ty: ignore[invalid-argument-type]
 
     # Auto-rotate x-axis labels if needed (not applicable for pie charts or world maps)
     if plot_type not in ["pie", "worldmap"]:
@@ -190,7 +194,7 @@ def _create_matplotlib_plot(  # noqa: C901
 
 def plot_to_bytes(df: pd.DataFrame, plot_type: str, **kwargs) -> bytes:
     """Generate a plot and return it as bytes."""
-    fig, _ = _create_matplotlib_plot(df, plot_type, **kwargs)
+    fig, _ = _create_plot(df, plot_type, **kwargs)
     buffer = io.BytesIO()
     fig.savefig(buffer, format="png", bbox_inches="tight")
     plt.close(fig)
@@ -200,7 +204,7 @@ def plot_to_bytes(df: pd.DataFrame, plot_type: str, **kwargs) -> bytes:
 
 def plot_and_show(df: pd.DataFrame, plot_type: str, **kwargs) -> None:
     """Generate a plot and display it."""
-    fig, _ = _create_matplotlib_plot(df, plot_type, **kwargs)
+    fig, _ = _create_plot(df, plot_type, **kwargs)
     plt.show()
     plt.close(fig)
 

tests/__init__.py

@@ -0,0 +1 @@
+# Tests package

tests/conftest.py

@@ -0,0 +1,9 @@
+"""Pytest configuration and shared fixtures."""
+
+import matplotlib
+
+
+def pytest_configure(config):
+    """Configure pytest with custom settings."""
+    # Use non-interactive backend for matplotlib to avoid GUI issues in tests
+    matplotlib.use("Agg")

tests/test_plot.py

@@ -0,0 +1,221 @@
+"""Tests for plotting functionality."""
+
+import matplotlib.pyplot as plt
+import pandas as pd
+import pytest
+from matplotlib.figure import Figure
+
+from plotting_mcp.plot import (
+    _auto_rotate_labels,
+    _create_pie_plot,
+    _create_plot,
+    plot_to_bytes,
+)
+
+
+class TestAutoRotateLabels:
+    """Test the _auto_rotate_labels function."""
+
+    def test_auto_rotate_labels_x_axis(self):
+        """Test auto rotation for x-axis labels."""
+        fig, ax = plt.subplots()
+        # Create long labels that should trigger rotation
+        long_labels = [f"Very Long Label {i}" for i in range(10)]
+        ax.set_xticks(range(len(long_labels)))
+        ax.set_xticklabels(long_labels)
+
+        _auto_rotate_labels(ax, axis="x")
+
+        # Check that rotation was applied
+        for label in ax.get_xticklabels():
+            assert label.get_rotation() == 90
+
+        plt.close(fig)
+
+    def test_auto_rotate_labels_y_axis(self):
+        """Test auto rotation for y-axis labels."""
+        fig, ax = plt.subplots()
+        # Create long labels that should trigger rotation
+        long_labels = [f"Very Long Label {i}" for i in range(10)]
+        ax.set_yticks(range(len(long_labels)))
+        ax.set_yticklabels(long_labels)
+
+        _auto_rotate_labels(ax, axis="y")
+
+        # Check that rotation was applied
+        for label in ax.get_yticklabels():
+            assert label.get_rotation() == 90
+
+        plt.close(fig)
+
+    def test_auto_rotate_labels_short_labels(self):
+        """Test that short labels don't get rotated."""
+        fig, ax = plt.subplots()
+        # Create short labels that shouldn't trigger rotation
+        short_labels = ["A", "B", "C", "D"]
+        ax.set_xticks(range(len(short_labels)))
+        ax.set_xticklabels(short_labels)
+
+        _auto_rotate_labels(ax, axis="x")
+
+        # Check that rotation was not applied (default rotation is 0)
+        for label in ax.get_xticklabels():
+            assert label.get_rotation() == 0
+
+        plt.close(fig)
+
+    def test_auto_rotate_labels_invalid_axis(self):
+        """Test that invalid axis raises ValueError."""
+        fig, ax = plt.subplots()
+
+        with pytest.raises(ValueError, match="Axis must be 'x' or 'y'"):
+            _auto_rotate_labels(ax, axis="z")
+
+        plt.close(fig)
+
+    def test_auto_rotate_labels_empty_labels(self):
+        """Test that empty labels don't cause errors."""
+        fig, ax = plt.subplots()
+        # No labels set
+
+        # Should not raise any errors
+        _auto_rotate_labels(ax, axis="x")
+        _auto_rotate_labels(ax, axis="y")
+
+        plt.close(fig)
+
+
+class TestCreatePiePlot:
+    """Test the _create_pie_plot function."""
+
+    def test_create_pie_plot_single_column_value_counts(self):
+        """Test pie plot with single column using value counts."""
+        df = pd.DataFrame({"category": ["A", "B", "A", "C", "B", "A"]})
+        fig, ax = plt.subplots()
+
+        _create_pie_plot(ax, df)
+
+        # Check that pie chart was created (wedges should exist)
+        assert len(ax.patches) > 0
+        plt.close(fig)
+
+    def test_create_pie_plot_two_columns(self):
+        """Test pie plot with two columns (labels and values)."""
+        df = pd.DataFrame({"category": ["A", "B", "C"], "values": [30, 45, 25]})
+        fig, ax = plt.subplots()
+
+        _create_pie_plot(ax, df)
+
+        # Check that pie chart was created with 3 wedges
+        assert len(ax.patches) == 3
+        plt.close(fig)
+
+    def test_create_pie_plot_too_many_columns(self):
+        """Test that pie plot with too many columns raises ValueError."""
+        df = pd.DataFrame({"col1": [1, 2, 3], "col2": [4, 5, 6], "col3": [7, 8, 9]})
+        fig, ax = plt.subplots()
+
+        with pytest.raises(ValueError, match="Pie chart requires either one column"):
+            _create_pie_plot(ax, df)
+
+        plt.close(fig)
+
+    def test_create_pie_plot_two_columns_with_labels_param(self):
+        """Test that pie plot with two columns rejects labels parameter."""
+        df = pd.DataFrame({"category": ["A", "B", "C"], "values": [30, 45, 25]})
+        fig, ax = plt.subplots()
+
+        with pytest.raises(ValueError, match="does not accept 'labels' parameter"):
+            _create_pie_plot(ax, df, labels=["X", "Y", "Z"])
+
+        plt.close(fig)
+
+
+class TestCreateMatplotlibPlot:
+    """Test the _create_matplotlib_plot function."""
+
+    def test_create_line_plot(self):
+        """Test creation of line plot."""
+        df = pd.DataFrame({"x": [1, 2, 3, 4, 5], "y": [2, 4, 6, 8, 10]})
+
+        fig, ax = _create_plot(df, "line", x="x", y="y")
+
+        assert isinstance(fig, Figure)
+        assert len(ax.lines) > 0  # Line plot should have lines
+        plt.close(fig)
+
+    def test_create_bar_plot(self):
+        """Test creation of bar plot."""
+        df = pd.DataFrame({"category": ["A", "B", "C", "D"], "values": [10, 15, 8, 12]})
+
+        fig, ax = _create_plot(df, "bar", x="category", y="values")
+
+        assert isinstance(fig, Figure)
+        assert len(ax.patches) > 0  # Bar plot should have bars (patches)
+        plt.close(fig)
+
+    def test_create_pie_plot(self):
+        """Test creation of pie plot through matplotlib function."""
+        df = pd.DataFrame({"category": ["A", "B", "C"], "values": [30, 45, 25]})
+
+        fig, ax = _create_plot(df, "pie")
+
+        assert isinstance(fig, Figure)
+        assert len(ax.patches) > 0  # Pie plot should have wedges (patches)
+        plt.close(fig)
+
+    def test_empty_dataframe_raises_error(self):
+        """Test that empty DataFrame raises ValueError."""
+        df = pd.DataFrame()
+
+        with pytest.raises(ValueError, match="CSV data is empty"):
+            _create_plot(df, "line")
+
+    def test_unsupported_plot_type_raises_error(self):
+        """Test that unsupported plot type raises ValueError."""
+        df = pd.DataFrame({"x": [1, 2, 3], "y": [4, 5, 6]})
+
+        with pytest.raises(ValueError, match="Unsupported plot type"):
+            _create_plot(df, "scatter3d")
+
+    def test_plot_with_title_and_labels(self):
+        """Test plot creation with title and axis labels."""
+        df = pd.DataFrame({"x": [1, 2, 3], "y": [4, 5, 6]})
+
+        fig, ax = _create_plot(
+            df, "line", x="x", y="y", title="Test Plot", xlabel="X Axis", ylabel="Y Axis"
+        )
+
+        assert ax.get_title() == "Test Plot"
+        assert ax.get_xlabel() == "X Axis"
+        assert ax.get_ylabel() == "Y Axis"
+        plt.close(fig)
+
+
+class TestPlotToBytes:
+    """Test the plot_to_bytes function."""
+
+    def test_plot_to_bytes_returns_bytes(self):
+        """Test that plot_to_bytes returns bytes."""
+        df = pd.DataFrame({"x": [1, 2, 3, 4, 5], "y": [2, 4, 6, 8, 10]})
+
+        result = plot_to_bytes(df, "line", x="x", y="y")
+
+        assert isinstance(result, bytes)
+        assert len(result) > 0
+        # Check PNG header
+        assert result.startswith(b"\x89PNG")
+
+    def test_plot_to_bytes_different_plot_types(self):
+        """Test plot_to_bytes with different plot types."""
+        df_line = pd.DataFrame({"x": [1, 2, 3], "y": [4, 5, 6]})
+        df_bar = pd.DataFrame({"cat": ["A", "B", "C"], "val": [1, 2, 3]})
+        df_pie = pd.DataFrame({"category": ["A", "B", "C"], "values": [30, 45, 25]})
+
+        line_bytes = plot_to_bytes(df_line, "line", x="x", y="y")
+        bar_bytes = plot_to_bytes(df_bar, "bar", x="cat", y="val")
+        pie_bytes = plot_to_bytes(df_pie, "pie")
+
+        assert all(isinstance(b, bytes) for b in [line_bytes, bar_bytes, pie_bytes])
+        assert all(len(b) > 0 for b in [line_bytes, bar_bytes, pie_bytes])
+        assert all(b.startswith(b"\x89PNG") for b in [line_bytes, bar_bytes, pie_bytes])