Feat/configurable timeout

README.md

@@ -74,6 +74,29 @@ wt_hist.get_historical_csv(
 )
 ```
 
+#### Tuning requests that time out
+
+Historical pulls are split into 30-day requests, each with a 60 s read timeout. If the API is slow to answer a large span, a request can time out and the error message will say so. Two options, which can be combined:
+
+```python
+from datetime import timedelta
+from watttime import WattTimeHistorical
+
+# allow each request more time: (connect, read) seconds
+wt_hist = WattTimeHistorical(username, password, timeout=(10, 300))
+
+# or make each request cover less time
+moers = wt_hist.get_historical_pandas(
+    start = '2022-01-01 00:00Z',
+    end = '2023-01-01 00:00Z',
+    region = 'CAISO_NORTH',
+    signal_type = 'co2_moer',
+    chunk_size = timedelta(days=10)
+)
+```
+
+The timeout applies to each attempt; the client retries a failed request up to three times before raising.
+
 You could also combine these classes to iterate through all regions where you have access to data:
 
 ```python

tests/test_sdk.py

@@ -1,3 +1,4 @@
+import inspect
 import unittest
 import unittest.mock as mock
 from unittest.mock import patch
@@ -15,6 +16,8 @@
 )
 from pathlib import Path
 import pytest
+import requests
+import urllib3.exceptions
 from shapely.geometry import shape, Polygon, MultiPolygon
 
 import pandas as pd
@@ -74,6 +77,119 @@ def test_login_sets_user_agent_header(self):
             self.base.headers["User-Agent"], f"watttime-python-sdk-{VERSION}"
         )
 
+    def _mock_login_response(self):
+        mock_rsp = mock.Mock()
+        mock_rsp.json.return_value = {"token": "fake-token"}
+        mock_rsp.raise_for_status.return_value = None
+        return mock_rsp
+
+    def _authed_base(self, **kwargs):
+        """A client with a fake, unexpired token so requests skip _login()."""
+        base = WattTimeBase(**kwargs)
+        base.token = "fake-token"
+        base.token_valid_until = datetime.now() + timedelta(minutes=30)
+        base.headers = {"Authorization": "Bearer fake-token"}
+        return base
+
+    def test_timeout_reaches_all_three_call_sites(self):
+        """The default (10, 60) is passed to login, the data request, and register."""
+        base = WattTimeBase()
+        self.assertEqual(base.timeout, (10, 60))
+
+        mock_rsp = mock.Mock()
+        mock_rsp.json.return_value = {"token": "fake-token", "data": [], "meta": {}}
+        mock_rsp.raise_for_status.return_value = None
+
+        with mock.patch.object(base.session, "get", return_value=mock_rsp) as mock_get:
+            base._login()
+            self.assertEqual(mock_get.call_args.kwargs["timeout"], (10, 60))
+
+            base._make_rate_limited_request("https://api.watttime.org/v3/test", {})
+            self.assertEqual(mock_get.call_args.kwargs["timeout"], (10, 60))
+
+        with mock.patch.object(
+            base.session, "post", return_value=mock_rsp
+        ) as mock_post:
+            base.register(email="someone@example.com")
+        self.assertEqual(mock_post.call_args.kwargs["timeout"], (10, 60))
+
+        base.session.close()
+
+    def test_timeout_accepts_tuple_and_scalar_and_none(self):
+        """Both `requests` timeout forms, and None, are stored and forwarded as given."""
+        for value in [(3, 5), 30, 12.5, None]:
+            with self.subTest(timeout=value):
+                base = WattTimeBase(timeout=value)
+                self.assertEqual(base.timeout, value)
+
+                with mock.patch.object(
+                    base.session, "get", return_value=self._mock_login_response()
+                ) as mock_get:
+                    base._login()
+
+                self.assertEqual(mock_get.call_args.kwargs["timeout"], value)
+                base.session.close()
+
+    def test_new_parameters_are_keyword_only(self):
+        """New optional parameters are keyword-only, so existing positional calls
+        are unaffected and the positional contract stays where it is today."""
+        keyword_only = inspect.Parameter.KEYWORD_ONLY
+        checks = [
+            (WattTimeBase.__init__, "timeout"),
+            (WattTimeHistorical.get_historical_jsons, "chunk_size"),
+            (WattTimeHistorical.get_historical_pandas, "chunk_size"),
+            (WattTimeHistorical.get_historical_csv, "chunk_size"),
+        ]
+        for func, name in checks:
+            with self.subTest(f"{func.__qualname__}({name})"):
+                param = inspect.signature(func).parameters[name]
+                self.assertIs(param.kind, keyword_only)
+
+        client = WattTimeHistorical(None, None, False, 10, 4)
+        self.assertEqual(client.worker_count, 4)
+        self.assertEqual(client.timeout, (10, 60))
+        client.session.close()
+
+        with self.assertRaises(TypeError):
+            WattTimeHistorical(None, None, False, 10, 4, (3, 5))
+
+    def test_timeout_error_message_points_at_knobs(self):
+        """A timeout is wrapped in RuntimeError with __cause__ preserved and a hint
+        naming both knobs -- including when urllib3's retry machinery has wrapped it
+        as ConnectionError -> MaxRetryError -> ReadTimeoutError, which is what the
+        session raises once its retries are exhausted."""
+        url = "https://api.watttime.org/v3/test"
+        reason = urllib3.exceptions.ReadTimeoutError(None, url, "Read timed out.")
+        exhausted = requests.exceptions.ConnectionError(
+            urllib3.exceptions.MaxRetryError(None, url, reason=reason)
+        )
+
+        for exc in [requests.exceptions.ReadTimeout("too slow"), exhausted]:
+            with self.subTest(exc=type(exc).__name__):
+                base = self._authed_base(timeout=(10, 60))
+                with mock.patch.object(base.session, "get", side_effect=exc):
+                    with self.assertRaises(RuntimeError) as ctx:
+                        base._make_rate_limited_request(url, {})
+
+                self.assertIs(ctx.exception.__cause__, exc)
+                message = str(ctx.exception)
+                self.assertIn("Hint", message)
+                self.assertIn("chunk_size", message)
+                self.assertIn("`timeout`", message)
+                self.assertIn("(10, 60)", message)
+                base.session.close()
+
+    def test_non_timeout_error_message_has_no_hint(self):
+        base = self._authed_base()
+        error = requests.exceptions.HTTPError("403 Client Error: Forbidden")
+
+        with mock.patch.object(base.session, "get", side_effect=error):
+            with self.assertRaises(RuntimeError) as ctx:
+                base._make_rate_limited_request("https://api.watttime.org/v3/test", {})
+
+        self.assertNotIn("Hint", str(ctx.exception))
+        base.session.close()
+
     @patch("time.sleep", return_value=None)
     def test_apply_rate_limit(self, mock_sleep):
         """Test _apply_rate_limit (single-threaded) triggers sleep when rate limit is exceeded."""
@@ -179,6 +295,34 @@ def test_get_chunks_rejects_inverted_span(self):
         with self.assertRaises(ValueError):
             self.base._get_chunks(start, end)
 
+    def test_get_chunks_honors_chunk_size(self):
+        start = datetime(2025, 1, 1, tzinfo=UTC)
+        end = datetime(2025, 1, 31, tzinfo=UTC)
+
+        chunks = self.base._get_chunks(start, end, chunk_size=timedelta(days=10))
+
+        self.assertEqual(len(chunks), 3)
+        self.assertEqual(chunks[0][0], start)
+        self.assertEqual(chunks[-1][1], end)
+        self.assertEqual(chunks[1][0] - chunks[0][1], timedelta(minutes=5))
+
+        # None means the 30-day default
+        self.assertEqual(
+            self.base._get_chunks(start, end, chunk_size=None),
+            self.base._get_chunks(start, end),
+        )
+
+    def test_get_chunks_rejects_chunk_size_at_or_below_trim(self):
+        # 5 minutes is trimmed from every chunk end but the last, so a chunk that
+        # short would be empty or inverted; zero or negative would never terminate.
+        start = datetime(2025, 1, 1, tzinfo=UTC)
+        end = datetime(2025, 1, 2, tzinfo=UTC)
+
+        for bad in [timedelta(minutes=5), timedelta(0), timedelta(minutes=-10)]:
+            with self.subTest(chunk_size=bad):
+                with self.assertRaises(ValueError):
+                    self.base._get_chunks(start, end, chunk_size=bad)
+
     @mock.patch("watttime.requests.Session.post", side_effect=mocked_register)
     def test_mock_register(self, mock_post):
         resp = self.base.register(email=os.getenv("WATTTIME_EMAIL"))
@@ -202,6 +346,38 @@ def setUp(self):
     def tearDown(self):
         self.historical.session.close()
 
+    def test_get_historical_jsons_chunk_size_controls_request_count(self):
+        """30 days at 10-day chunks is three requests instead of one."""
+        fake_response = {"data": [], "meta": {"model": {"date": "2024-01-01"}}}
+
+        with mock.patch.object(
+            self.historical, "_fetch_data", return_value=[fake_response]
+        ) as mock_fetch:
+            self.historical.get_historical_jsons(
+                "2025-01-01 00:00Z",
+                "2025-01-31 00:00Z",
+                REGION,
+                chunk_size=timedelta(days=10),
+            )
+
+        param_chunks = mock_fetch.call_args.args[1]
+        self.assertEqual(len(param_chunks), 3)
+
+    def test_get_historical_pandas_forwards_chunk_size(self):
+        fake_response = {
+            "data": [{"point_time": "2025-01-01T00:00:00+00:00", "value": 1.0}],
+            "meta": {},
+        }
+
+        with mock.patch.object(
+            self.historical, "get_historical_jsons", return_value=[fake_response]
+        ) as mock_jsons:
+            self.historical.get_historical_pandas(
+                "2025-01-01", "2025-01-02", REGION, chunk_size=timedelta(days=1)
+            )
+
+        self.assertEqual(mock_jsons.call_args.kwargs["chunk_size"], timedelta(days=1))
+
     def test_get_historical_jsons_3_months(self):
         start = "2025-01-01 00:00Z"
         end = "2025-03-31 00:00Z"
@@ -289,6 +465,24 @@ def test_get_historical_csv(self):
         assert fp.exists()
         fp.unlink()
 
+    def test_get_historical_csv_include_imputed(self):
+        start = parse("2025-01-01 00:00Z")
+        end = parse("2025-01-02 00:00Z")
+        self.historical.get_historical_csv(
+            start, end, REGION, include_imputed_marker=True
+        )
+
+        fp = (
+            Path.home()
+            / "watttime_historical_csvs"
+            / f"{REGION}_co2_moer_{start.date()}_{end.date()}.csv"
+        )
+        assert fp.exists()
+        df = pd.read_csv(fp)
+        self.assertIn("imputed_data_used", df.columns)
+        self.assertNotIn("meta", df.columns)
+        fp.unlink()
+
     def test_multi_model_range(self):
         """If model is not specified, we should only return the most recent model data"""
         myaccess = WattTimeMyAccess()