Merge pull request #2 from PR0M3TH3AN/codex/modify-plot_lagged-to-return-matplotlib-figure

Return Matplotlib figure and test plot

AGENTS.md

@@ -69,6 +69,7 @@ Future extensions will add additional custom chart scripts under `scripts/`, all
 
 * **Automated tests**: Add pytest modules under `tests/` following the pattern `test_*.py`.
 * **Data mocking**: For new chart scripts, mock FRED downloads via monkeypatching `pandas_datareader.data.DataReader`.
+* **Figure tests**: `tests/test_plot.py` ensures `plot_lagged` returns a figure with two axes and at least one line.
 
 ## Extending with New Charts
 

scripts/lagged_oil_unrate_chart_styled.py

@@ -45,20 +45,18 @@ def fetch_series(
         )
 
     with sqlite3.connect(db_path) as conn:
-        unrate = (
-            pd.read_sql(
-                "SELECT date, UNRATE AS value FROM UNRATE", conn,
-                parse_dates=["date"], index_col="date"
-            )
-            .loc[start:end]
-        )
-        oil = (
-            pd.read_sql(
-                "SELECT date, DCOILWTICO AS value FROM DCOILWTICO", conn,
-                parse_dates=["date"], index_col="date"
-            )
-            .loc[start:end]
-        )
+        unrate = pd.read_sql(
+            "SELECT date, UNRATE AS value FROM UNRATE",
+            conn,
+            parse_dates=["date"],
+            index_col="date",
+        ).loc[start:end]
+        oil = pd.read_sql(
+            "SELECT date, DCOILWTICO AS value FROM DCOILWTICO",
+            conn,
+            parse_dates=["date"],
+            index_col="date",
+        ).loc[start:end]
 
     # Resample WTI to month-end average
     oil_monthly = oil.resample("M").mean()
@@ -108,17 +106,22 @@ def plot_lagged(
     _start_date: datetime,
     _end_date: datetime,
     extend_years: int,
-) -> None:
+) -> plt.Figure:
     """
     Render the dual-axis chart with fixed scales and log ticks on oil.
     Trims both series to their overlapping date range before shifting.
+
+    Returns the created :class:`matplotlib.figure.Figure` so callers can
+    decide how to display or save it.
     """
     # Find overlapping date range
     common_start = max(unrate.index.min(), oil.index.min())
     common_end = min(unrate.index.max(), oil.index.max())
 
     print(f"⚠️ Using overlapping range: {common_start.date()} to {common_end.date()}")
-    print(f"  - UNRATE range: {unrate.index.min().date()} to {unrate.index.max().date()}")
+    print(
+        f"  - UNRATE range: {unrate.index.min().date()} to {unrate.index.max().date()}"
+    )
     print(f"  - OIL range: {oil.index.min().date()} to {oil.index.max().date()}")
 
     # Trim series to common range
@@ -208,7 +211,7 @@ def plot_lagged(
     )
 
     fig.tight_layout(rect=[0, 0.03, 1, 0.95])
-    plt.show()
+    return fig
 
 
 # ──────────────────────────────────────────────────────────────────────────
@@ -254,7 +257,8 @@ def main() -> None:
 
     unrate, oil = fetch_series(start_dt, end_dt, args.db)
     validate_series(unrate, oil)
-    plot_lagged(unrate, oil, args.offset, start_dt, end_dt, args.extend_years)
+    fig = plot_lagged(unrate, oil, args.offset, start_dt, end_dt, args.extend_years)
+    fig.show()
 
 
 if __name__ == "__main__":

tests/test_plot.py

@@ -0,0 +1,22 @@
+import matplotlib
+
+matplotlib.use("Agg")
+
+import os
+import sys
+import pandas as pd
+
+sys.path.append(os.path.abspath("."))
+
+from scripts.lagged_oil_unrate_chart_styled import plot_lagged
+
+
+def test_plot_lagged_returns_figure():
+    dates = pd.date_range("2000-01-01", periods=5, freq="M")
+    unrate = pd.DataFrame({"value": [5, 6, 7, 5, 4]}, index=dates)
+    oil = pd.DataFrame({"value": [20, 25, 30, 35, 40]}, index=dates)
+
+    fig = plot_lagged(unrate, oil, 1, dates[0], dates[-1], 0)
+
+    assert len(fig.axes) == 2
+    assert any(ax.get_lines() for ax in fig.axes)