Skip to content

Commit 6ebf3a9

Browse files
authored
Add design docs (#109)
1 parent 31abf9b commit 6ebf3a9

2 files changed

Lines changed: 276 additions & 0 deletions

File tree

docs/design_docs/data_fetcher.md

Lines changed: 127 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,127 @@
1+
# Design Document: Data Fetcher Implementation
2+
3+
This document provides a guide for implementing a new river gauge data fetcher for the RivRetrieve library.
4+
5+
## Goal
6+
7+
The goal is to create a standardized interface for downloading and parsing streamflow, water level and any other river measurement data from various national and regional data providers.
8+
9+
## Architecture
10+
11+
Every data fetcher must inherit from the `RiverDataFetcher` abstract base class defined in `rivretrieve.base`.
12+
13+
### Key Components
14+
15+
#### 1. Variable Definitions (`rivretrieve.constants`)
16+
17+
All fetchers must use the standardized variable names defined in `rivretrieve.constants`. Common variables include:
18+
19+
- `DISCHARGE_DAILY_MEAN`
20+
- `DISCHARGE_INSTANT`
21+
- `STAGE_DAILY_MEAN`
22+
- `STAGE_INSTANT`
23+
24+
If data providers include variables that are not yet defined in `rivretrieve.constants`, we define
25+
a new name there, then again used the globally defined constant.
26+
27+
Output data must be converted to SI units, e.g.
28+
- Discharge: Cubic meters per second (m³/s)
29+
- Stage: Meters (m)
30+
31+
#### 2. Class Structure
32+
33+
A new fetcher class (e.g., `USAFetcher`) should be implemented in its own file (e.g., `rivretrieve/usa.py`).
34+
35+
```python
36+
from typing import Optional
37+
import pandas as pd
38+
from . import base, constants, utils
39+
40+
class MyCountryFetcher(base.RiverDataFetcher):
41+
# Implementation details...
42+
```
43+
44+
#### 3. Common Utility Functions (`rivretrieve.utils`)
45+
46+
- **`format_start_date(start_date)` / `format_end_date(end_date)`**: Standardizes date strings to 'YYYY-MM-DD'.
47+
- **`requests_retry_session()`**: Returns a `requests.Session` object with built-in retry logic for handling transient network errors.
48+
- **`load_cached_metadata_csv(country_code)`**: Loads the site metadata from `rivretrieve/cached_site_data/{country_code}_sites.csv`.
49+
50+
#### 4. Mandatory Methods
51+
52+
- **`get_available_variables() -> tuple[str, ...]`**:
53+
Returns a tuple of the `constants` supported by this fetcher.
54+
55+
- **`get_cached_metadata() -> pd.DataFrame`**:
56+
Retrieves available gauge IDs and metadata from a cached CSV file. Use `utils.load_cached_metadata_csv("country_name")`.
57+
58+
- **`_download_data(gauge_id, variable, start_date, end_date) -> any`**:
59+
Handles the low-level data retrieval (e.g., via `requests` or a provider-specific library).
60+
- `start_date` and `end_date` are strings in 'YYYY-MM-DD' format.
61+
- Returns raw data (e.g., a `pd.DataFrame`, `dict`, or `str`).
62+
63+
- **`_parse_data(gauge_id, raw_data, variable) -> pd.DataFrame`**:
64+
Parses the raw data into a standardized `pd.DataFrame`.
65+
- Index: `pd.DatetimeIndex` named `constants.TIME_INDEX`.
66+
- Column: A single column named after the `variable`.
67+
- Handles unit conversions to SI.
68+
- Handles missing data (NaN). Important: Different countries might use different constants to
69+
indicate missing data (e.g. sometimes negativ values like `-999`, sometimes strings `MISSING`, `LUECKE`). We always want to convert these country specific constants to `np.nan`.
70+
71+
- **`get_data(gauge_id, variable, start_date, end_date) -> pd.DataFrame`**:
72+
The main entry point for users. It should:
73+
1. Format dates using `utils.format_start_date` and `utils.format_end_date`.
74+
2. Validate the `variable`.
75+
3. Call `_download_data` and `_parse_data`.
76+
4. Return the standardized `pd.DataFrame`.
77+
78+
#### 5. Metadata Handling
79+
80+
Metadata should be cached as a CSV file in `rivretrieve/cached_site_data/`. The CSV should use standard column names from `constants.py` for commonly used information:
81+
- `GAUGE_ID` (index)
82+
- `STATION_NAME`
83+
- `RIVER`
84+
- `LATITUDE`
85+
- `LONGITUDE`
86+
- `ALTITUDE`
87+
- `AREA`
88+
- `COUNTRY`
89+
90+
However, the metadata doesn't have to be restricted to these columns and can include any additional
91+
column with it's original column name.
92+
93+
#### 6. Optional Methods
94+
95+
- **`get_metadata(self) -> pd.DataFrame`**:
96+
Downloads and parses site metadata directly from the data provider. If a live metadata endpoint is available, this method should download the raw data, rename the columns to the standard `constants`, add `constants.COUNTRY` and `constants.SOURCE` where appropriate, ensure coordinate types are correctly converted, and return a DataFrame indexed by `constants.GAUGE_ID`.
97+
98+
## Implementation Steps
99+
100+
1. **Identify the Data Source**: Determine the provider's API or data download URL.
101+
2. **Define Supported Variables**: Map the provider's variables to `rivretrieve.constants`.
102+
3. **Implement `_download_data`**: Use `requests` or other tools to fetch raw data.
103+
4. **Implement `_parse_data`**: Convert the raw format to the standardized `pd.DataFrame`.
104+
5. **Create Metadata**: Prepare the `cached_site_data/country.csv` file.
105+
106+
107+
### 6. Verification
108+
109+
- **Example Script**: Add a script to `examples/` (e.g., `download_mycountry_data.py`) that demonstrates using the new fetcher for a single gauge and plots the result.
110+
- **Unit Tests**: **Crucial Step**. You must create a corresponding test file (e.g. `tests/test_country.py`).
111+
- **See the full testing guide in [data_fetcher_test.md](data_fetcher_test.md) for detailed instructions.**
112+
- **The Golden Rule**: 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. The mocked data must be a **real, raw payload** obtained from the API and stored in `tests/test_data/`.
113+
- Use `pandas.testing.assert_frame_equal` to compare the fetcher's output against a known `expected_df`.
114+
115+
## Best Practices
116+
117+
- **Standardized Empty DataFrames**: If an API request fails or no data is found, always catch exceptions and return an empty DataFrame with the correct columns: `pd.DataFrame(columns=[constants.TIME_INDEX, variable])`. Do not return `None`.
118+
- **Date Filtering**: APIs frequently return data in whole months or years. Make sure the final return in `get_data()` perfectly filters the DataFrame to exactly match the requested `start_date` and `end_date` using `df[(df.index >= start_date_dt) & (df.index <= end_date_dt)]`.
119+
- **Authentication & Credentials**: If the provider requires an API key or password, use `python-dotenv` and load credentials from a `.env` file (e.g., `os.environ.get("MY_API_KEY")`). Ensure `__init__` can optionally accept these credentials directly as kwargs.
120+
- **Pagination & Chunking**: When fetching large time ranges, chunk the requests (e.g., by year or month) within `_download_data` to prevent timeout or payload size errors.
121+
- **API Limits & Throttling**: Some APIs have strict request limits. Implement proper throttling (e.g., with `time.sleep()`) and handle `HTTP 429` appropriately to be respectful of external servers.
122+
- **Bulk Downloads & Caching**: For providers without a robust time-series API, a common architectural pattern is to download bulk datasets (e.g., a large zip file) on the first request, save it to `rivretrieve/data/`, cache the processed data locally (e.g., as `.zarr` or `.sqlite3`), and serve subsequent queries directly from this local cache.
123+
- **Class Docstrings**: Ensure the fetcher class has a docstring specifying the "Data Source:" (with a URL) and "Supported Variables:" (listing the `constants` used).
124+
- Use `logging` for errors and warnings.
125+
- Use `pd.to_numeric(..., errors="coerce")` to handle malformed data gracefully.
126+
- Ensure all datetime objects are timezone-naive or consistently handled (prefer UTC).
127+
- Avoid dropping columns that are not explicitly renamed during metadata parsing.
Lines changed: 149 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,149 @@
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

Comments
 (0)