|
| 1 | +# Design Document: Data Fetcher Unit Tests |
| 2 | + |
| 3 | +This document provides a comprehensive guide for writing unit tests for new river gauge data fetchers in the RivRetrieve library. |
| 4 | + |
| 5 | +## The Golden Rule of RivRetrieve Testing |
| 6 | + |
| 7 | +**Each unit test must ONLY mock the call to the external data provider. Everything else from our code (parsing, date formatting, unit conversions) MUST be tested.** |
| 8 | + |
| 9 | +The mocked data must be a **real, raw payload** obtained from the API. We do not invent mock data structures; we capture real API responses and use them to ensure our parsing logic works against the actual data formats provided by the sources. However, it is enough to test against a short time period |
| 10 | +of a few days. |
| 11 | + |
| 12 | +## Directory Structure |
| 13 | + |
| 14 | +- **Test File:** `tests/test_<country>.py` (e.g., `tests/test_brazil.py`) |
| 15 | +- **Test Data:** `tests/test_data/<country>_<gauge_id>_<variable>_<date>.<ext>` (e.g., `tests/test_data/uk_nrfa_1001_discharge_20220101.json`) |
| 16 | + |
| 17 | +## 1. Obtaining Test Data |
| 18 | + |
| 19 | +Before writing the test, you need real payloads. |
| 20 | +1. Temporarily add print statements or a debugger to your fetcher's `_download_data` method just before it parses the raw response. |
| 21 | +2. Run your fetcher using a script in `examples/` for a short time range (e.g., 3-5 days). |
| 22 | +3. Save the exact raw response (JSON, CSV, HTML, XML, or binary) to a file in the `tests/test_data/` directory. |
| 23 | +4. *Exception*: If the payload is extremely small (e.g., a simple JSON dict with a few keys), you can define it directly in the test file as a Python dictionary. |
| 24 | + |
| 25 | +## 2. Test Class Structure |
| 26 | + |
| 27 | +All tests should inherit from `unittest.TestCase`. |
| 28 | + |
| 29 | +### `setUp` Method |
| 30 | +Use the `setUp` method to initialize your fetcher and define the path to your test data. |
| 31 | + |
| 32 | +```python |
| 33 | +import os |
| 34 | +import json |
| 35 | +import unittest |
| 36 | +from unittest.mock import MagicMock, patch |
| 37 | +import pandas as pd |
| 38 | +from pandas.testing import assert_frame_equal |
| 39 | +from rivretrieve import MyCountryFetcher, constants |
| 40 | + |
| 41 | +class TestMyCountryFetcher(unittest.TestCase): |
| 42 | + def setUp(self): |
| 43 | + self.fetcher = MyCountryFetcher() |
| 44 | + self.test_data_dir = os.path.join(os.path.dirname(__file__), "test_data") |
| 45 | + |
| 46 | + def load_sample_data(self, filename): |
| 47 | + with open(os.path.join(self.test_data_dir, filename), "r", encoding="utf-8") as f: |
| 48 | + return f.read() |
| 49 | + |
| 50 | + def load_sample_json(self, filename): |
| 51 | + with open(os.path.join(self.test_data_dir, filename), "r", encoding="utf-8") as f: |
| 52 | + return json.load(f) |
| 53 | +``` |
| 54 | + |
| 55 | +## 3. Mocking Strategies |
| 56 | + |
| 57 | +You must mock the boundary where our code leaves the system. In 95% of cases, this is the `requests` library. |
| 58 | + |
| 59 | +### Mocking `requests_retry_session` |
| 60 | +If your fetcher uses `utils.requests_retry_session().get(...)`: |
| 61 | + |
| 62 | +```python |
| 63 | + @patch("rivretrieve.utils.requests_retry_session") |
| 64 | + def test_get_data_discharge(self, mock_requests_session): |
| 65 | + mock_session = MagicMock() |
| 66 | + mock_requests_session.return_value = mock_session |
| 67 | + |
| 68 | + mock_response = MagicMock() |
| 69 | + mock_response.text = self.load_sample_data("mycountry_sample.csv") |
| 70 | + # OR: mock_response.json.return_value = self.load_sample_json("mycountry_sample.json") |
| 71 | + mock_response.raise_for_status = MagicMock() |
| 72 | + |
| 73 | + mock_session.get.return_value = mock_response |
| 74 | + |
| 75 | + # ... proceed with calling fetcher.get_data(...) |
| 76 | +``` |
| 77 | + |
| 78 | +### Mocking Multiple Sequential API Calls |
| 79 | +If the fetcher needs to call multiple endpoints (e.g., one for a token/metadata, one for the actual data), use `side_effect`: |
| 80 | + |
| 81 | +```python |
| 82 | + def mock_get_side_effect(url, *args, **kwargs): |
| 83 | + mock_response = MagicMock() |
| 84 | + if "metadata_endpoint" in url: |
| 85 | + mock_response.json.return_value = self.load_sample_json("meta.json") |
| 86 | + elif "data_endpoint" in url: |
| 87 | + mock_response.json.return_value = self.load_sample_json("data.json") |
| 88 | + mock_response.raise_for_status = MagicMock() |
| 89 | + return mock_response |
| 90 | + |
| 91 | + mock_session.get.side_effect = mock_get_side_effect |
| 92 | +``` |
| 93 | + |
| 94 | +### Mocking External Libraries |
| 95 | +If the fetcher uses a dedicated external client library (e.g., `dataretrieval` for USA), mock the library's function: |
| 96 | + |
| 97 | +```python |
| 98 | + @patch("dataretrieval.nwis.get_dv") |
| 99 | + def test_get_data_discharge(self, mock_get_dv): |
| 100 | + mock_get_dv.return_value = (self.load_sample_csv_as_df(), MagicMock()) |
| 101 | +``` |
| 102 | + |
| 103 | +## 4. Assertions and Validation |
| 104 | + |
| 105 | +Your test must execute `get_data()` and validate the returned DataFrame against an expected DataFrame constructed manually in the test. |
| 106 | + |
| 107 | +```python |
| 108 | + gauge_id = "12345" |
| 109 | + variable = constants.DISCHARGE_DAILY_MEAN |
| 110 | + start_date = "2020-01-01" |
| 111 | + end_date = "2020-01-03" |
| 112 | + |
| 113 | + result_df = self.fetcher.get_data(gauge_id, variable, start_date, end_date) |
| 114 | + |
| 115 | + # Build the exact expected output. |
| 116 | + # Make sure to apply any unit conversions here that the fetcher should have done! |
| 117 | + expected_dates = pd.to_datetime(["2020-01-01", "2020-01-02", "2020-01-03"]) |
| 118 | + expected_values = [10.5, 11.2, 9.8] # Already converted to SI units (m³/s) |
| 119 | + |
| 120 | + expected_data = { |
| 121 | + constants.TIME_INDEX: expected_dates, |
| 122 | + variable: expected_values, |
| 123 | + } |
| 124 | + expected_df = pd.DataFrame(expected_data).set_index(constants.TIME_INDEX) |
| 125 | + |
| 126 | + # Assert DataFrame matches perfectly |
| 127 | + assert_frame_equal(result_df, expected_df, check_dtype=False) |
| 128 | + |
| 129 | + # Assert the mocked API was called with the correct parameters |
| 130 | + mock_session.get.assert_called_once() |
| 131 | + args, kwargs = mock_session.get.call_args |
| 132 | + self.assertIn("12345", args[0] if args else kwargs.get("url", "")) |
| 133 | +``` |
| 134 | + |
| 135 | +## 5. Testing Metadata (`get_metadata`) |
| 136 | + |
| 137 | +If your fetcher implements the optional `get_metadata()` method, write a test for it: |
| 138 | +1. Save the raw metadata payload. |
| 139 | +2. Mock the request. |
| 140 | +3. Assert that the resulting DataFrame has the index named `constants.GAUGE_ID`. |
| 141 | +4. Assert that standard columns like `constants.STATION_NAME`, `constants.LATITUDE`, `constants.LONGITUDE`, etc., are present and correctly mapped. |
| 142 | + |
| 143 | +## Summary Checklist |
| 144 | +- [ ] Named `test_<country>.py` |
| 145 | +- [ ] Raw test payload saved in `tests/test_data/` |
| 146 | +- [ ] Only the HTTP call or external library call is mocked |
| 147 | +- [ ] DataFrame is compared using `assert_frame_equal` |
| 148 | +- [ ] Mock call arguments are verified (`assert_called_once_with`, etc.) |
| 149 | +- [ ] Target variables are tested independently (e.g., test Discharge, test Stage) |
0 commit comments