diff --git a/AGENTS.md b/AGENTS.md index 4ac6a91..298bee1 100644 --- a/AGENTS.md +++ b/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 diff --git a/scripts/lagged_oil_unrate_chart_styled.py b/scripts/lagged_oil_unrate_chart_styled.py index ef9c3ba..95978e0 100644 --- a/scripts/lagged_oil_unrate_chart_styled.py +++ b/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__": diff --git a/tests/test_plot.py b/tests/test_plot.py new file mode 100644 index 0000000..082496b --- /dev/null +++ b/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)