diff --git a/tests/test_sdk.py b/tests/test_sdk.py index 9ce4b77..55bc913 100644 --- a/tests/test_sdk.py +++ b/tests/test_sdk.py @@ -157,7 +157,7 @@ def test_username_arg_warns_when_env_user_already_set(self, mock_warning): class TestWattTimeHistorical(unittest.TestCase): def setUp(self): - self.historical = WattTimeHistorical(rate_limit=1) + self.historical = WattTimeHistorical() def tearDown(self): self.historical.session.close() @@ -208,6 +208,17 @@ def test_get_historical_pandas(self): self.assertIn("point_time", df.columns) self.assertIn("value", df.columns) + def test_get_historical_pandas_include_imputed(self): + start = datetime.now() - timedelta(days=7) + end = datetime.now() + df = self.historical.get_historical_pandas(start, end, REGION, include_imputed_marker=True) + + self.assertIsInstance(df, pd.DataFrame) + self.assertGreaterEqual(len(df), 1) + self.assertIn("point_time", df.columns) + self.assertIn("value", df.columns) + self.assertIn("imputed_data_used", df.columns) + def test_get_historical_pandas_meta(self): start = datetime.now() - timedelta(days=7) end = datetime.now() @@ -274,6 +285,9 @@ def setUp(self): def tearDown(self): self.historical.session.close() + def tearDown(self): + self.historical.session.close() + def test_get_historical_jsons_3_months_multithreaded(self): start = "2026-01-01 00:00Z" end = "2026-03-31 00:00Z" @@ -356,6 +370,9 @@ def setUp(self): def tearDown(self): self.forecast.session.close() + def tearDown(self): + self.forecast.session.close() + def test_get_current_json(self): json = self.forecast.get_forecast_json(region=REGION) @@ -459,6 +476,9 @@ def setUp(self): def tearDown(self): self.forecast.session.close() + def tearDown(self): + self.forecast.session.close() + def test_historical_forecast_jsons_multithreaded(self): start = "2026-01-01 00:00Z" end = "2026-01-14 00:00Z" @@ -532,75 +552,6 @@ def test_my_access_in_geojson(self): set(maps_regions) - set(access_regions) == set() ), f"Extra regions in geojson for {signal_type}: {set(maps_regions) - set(access_regions)}" - def test_ccw(self): - moer = self.maps.get_maps_json(signal_type="co2_moer") - - def _is_ccw(geometry): - if isinstance(geometry, Polygon): - return geometry.exterior.is_ccw - elif isinstance(geometry, MultiPolygon): - return all(poly.exterior.is_ccw for poly in geometry.geoms) - return True - - bad = [ - f["properties"]["region_full_name"] - for f in moer["features"] - if not _is_ccw(shape(f["geometry"])) - ] - assert len(bad) == 0, f"Non-CCW geometries: {bad}" - - -class TestWattTimeMarginalFuelMix(unittest.TestCase): - def setUp(self): - self.fuel_mix = WattTimeMarginalFuelMix() - - def test_fuel_mix_jsons(self): - start = "2026-01-01 00:00Z" - end = "2026-01-07 00:00Z" - fm = self.fuel_mix.get_fuel_mix_jsons(start=start, end=end, region=REGION) - - self.assertIsInstance(fm, list) - self.assertIsInstance(fm[0], dict) - self.assertIn("data", fm[0]) - self.assertIn("meta", fm[0]) - self.assertIsInstance(fm[0]["data"], list) - self.assertIsInstance(fm[0]["data"][0], dict) - self.assertIn("point_time", fm[0]["data"][0]) - self.assertIn("values", fm[0]["data"][0]) - self.assertIsInstance(fm[0]["data"][0]["values"], list) - self.assertIsInstance(fm[0]["data"][0]["values"][0], dict) - self.assertIn("fuel_type", fm[0]["data"][0]["values"][0]) - self.assertIn("value", fm[0]["data"][0]["values"][0]) - self.assertGreaterEqual(len(fm[0]["data"]), 12 * 24 * 5) - self.assertIn("warnings", fm[0]["meta"]) - self.assertIn("model", fm[0]["meta"]) - self.assertEqual("marginal_fuel_mix", fm[0]["meta"]["signal_type"]) - self.assertEqual("proportion", fm[0]["meta"]["units"]) - self.assertEqual(REGION, fm[0]["meta"]["region"]) - - def test_fuel_mix_pandas(self): - start = "2026-01-01 00:00Z" - end = "2026-01-07 00:00Z" - df = self.fuel_mix.get_fuel_mix_pandas(start=start, end=end, region=REGION) - self.assertIsInstance(df, pd.DataFrame) - self.assertGreaterEqual(len(df), 12 * 24 * 5) - self.assertIn("gas", df.columns) - self.assertEqual("point_time", df.index.name) - assert df.index.is_monotonic_increasing - assert all(df.sum(axis="columns") == 1.0) - - @patch.object(WattTimeMarginalFuelMix, "_fetch_data", side_effect=RuntimeError("403 Client Error: Forbidden")) - def test_get_fuel_mix_jsons_handles_403(self, mock_fetch_data): - start = "2026-01-01 00:00Z" - end = "2026-01-07 00:00Z" - - result = self.fuel_mix.get_fuel_mix_jsons(start=start, end=end, region=REGION) - - self.assertEqual(result, []) - mock_fetch_data.assert_called_once() - # TODO: test for logging here once we use log rather than print in api.py - - if __name__ == "__main__": unittest.main() diff --git a/watttime/api.py b/watttime/api.py index 199e1cc..a583356 100644 --- a/watttime/api.py +++ b/watttime/api.py @@ -113,6 +113,28 @@ def __init__( self.session.mount("http://", adapter) self.session.mount("https://", adapter) + @property + def password(self): + password = os.getenv("WATTTIME_PASSWORD") + if not password: + raise ValueError( + "WATTTIME_PASSWORD env variable is not set." + + "Please set this variable, or pass in a password upon initialization," + + "which will store it as a variable only for the current session" + ) + return password + + @property + def username(self): + username = os.getenv("WATTTIME_USER") + if not username: + raise ValueError( + "WATTTIME_USER env variable is not set." + + "Please set this variable, or pass in a username upon initialization," + + "which will store it as a variable only for the current session" + ) + return username + def _login(self): """ Login to the WattTime API, which provides a JWT valid for 30 minutes @@ -271,7 +293,7 @@ def _make_rate_limited_request(self, url: str, params: Dict[str, Any]) -> Dict: self._apply_rate_limit(ts) try: - rsp = self.session.get(url, headers=self.headers, params=params, timeout=60) + rsp = self.session.get(url, headers=self.headers, params=params) rsp.raise_for_status() j = rsp.json() except requests.exceptions.RequestException as e: @@ -293,6 +315,8 @@ def _make_rate_limited_request(self, url: str, params: Dict[str, Any]) -> Dict: self._last_request_meta = meta + self._last_request_meta = meta + return j def _apply_rate_limit(self, ts: float): @@ -370,6 +394,7 @@ def get_historical_jsons( Literal["co2_moer", "co2_aoer", "health_damage"] ] = "co2_moer", model: Optional[Union[str, date]] = None, + include_imputed_marker: bool = False, ) -> List[dict]: """ Base function to scrape historical data, returning a list of .json responses. @@ -391,6 +416,9 @@ def get_historical_jsons( url = "{}/v3/historical".format(self.url_base) params = {"region": region, "signal_type": signal_type} + if include_imputed_marker: + params["include_imputed_marker"] = "true" + start, end = self._parse_dates(start, end) chunks = self._get_chunks(start, end) @@ -421,6 +449,7 @@ def get_historical_pandas( ] = "co2_moer", model: Optional[Union[str, date]] = None, include_meta: bool = False, + include_imputed_marker: bool = False, ): """ Return a pd.DataFrame with point_time, and values. @@ -434,7 +463,9 @@ def get_historical_pandas( Returns: pd.DataFrame: _description_ """ - responses = self.get_historical_jsons(start, end, region, signal_type, model) + responses = self.get_historical_jsons( + start, end, region, signal_type, model, include_imputed_marker + ) df = pd.json_normalize( responses, record_path="data", meta=["meta"] if include_meta else [] ) @@ -452,6 +483,7 @@ def get_historical_csv( Literal["co2_moer", "co2_aoer", "health_damage"] ] = "co2_moer", model: Optional[Union[str, date]] = None, + include_imputed_marker: bool = False, ): """ Retrieves historical data from a specified start date to an end date and saves it as a CSV file. @@ -467,7 +499,9 @@ def get_historical_csv( Returns: None, results are saved to a csv file in the user's home directory. """ - df = self.get_historical_pandas(start, end, region, signal_type, model) + df = self.get_historical_pandas( + start, end, region, signal_type, model, include_imputed_marker + ) out_dir = Path.home() / "watttime_historical_csvs" out_dir.mkdir(exist_ok=True)