Skip to content

Commit a0dbe24

Browse files
authored
Merge pull request #33 from praekeltfoundation/retry-turn-calls
Retry turn calls
2 parents e4c0920 + 75b39e1 commit a0dbe24

File tree

12 files changed

+843
-415
lines changed

12 files changed

+843
-415
lines changed

README.md

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -122,3 +122,13 @@ pyS3.s3.get_filenames(bucket=bucket, prefix=prefix)
122122
#### Running tests
123123
1. `uv sync --dev --extra polars`
124124
2. `uv run pytest -vv --cov`
125+
126+
## Release
127+
To release a new version of the package:
128+
129+
1. Update the version number in `pyproject.toml`
130+
1. Run `uv lock`
131+
1. Post the changes to Slack for approval
132+
1. Once approved, push to main
133+
1. Repeat the above post-release, incrementing and adding `.dev0` to the version number.
134+

pyproject.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ dependencies = [
1313
"awswrangler>=3.7.3",
1414
"boto3>=1.34.103",
1515
"httpx>=0.27.0",
16+
"httpx-retries>=0.4.2",
1617
"pandas>=2.2.2",
1718
"types-tqdm>=4.66.0.20240417",
1819
]
Lines changed: 3 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,4 @@
1-
from httpx import Client
2-
3-
from .. import config_from_env
4-
5-
API_KEY = config_from_env("FLOW_RESULTS_API_KEY")
6-
BASE_URL = config_from_env("FLOW_RESULTS_API_BASE_URL")
7-
8-
9-
headers = {"Authorization": f"Token {API_KEY}"}
10-
11-
client: Client = Client(base_url=BASE_URL, headers=headers)
12-
1+
from .client import get_client, make_client
132
from .main import pyFlows as pyFlows
3+
4+
__all__ = ["get_client", "make_client", "pyFlows"]
Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,25 @@
1+
from functools import lru_cache
2+
3+
from httpx import Client
4+
from httpx_retries import RetryTransport
5+
6+
from .. import config_from_env
7+
8+
API_KEY = config_from_env("FLOW_RESULTS_API_KEY")
9+
BASE_URL = config_from_env("FLOW_RESULTS_API_BASE_URL")
10+
11+
headers = {"Authorization": f"Token {API_KEY}"}
12+
13+
14+
def make_client() -> Client:
15+
return Client(
16+
base_url=BASE_URL,
17+
headers=headers,
18+
timeout=30.0,
19+
transport=RetryTransport(),
20+
)
21+
22+
23+
@lru_cache(maxsize=1)
24+
def get_client() -> Client:
25+
return make_client()

rdw_ingestion_tools/api/flow_results/extensions/httpx.py

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,6 @@ def get_ids(client: Client, **kwargs: str | int) -> Iterator[str]:
1010
Results API.
1111
1212
"""
13-
1413
params = {**kwargs}
1514
url = ""
1615

@@ -28,7 +27,6 @@ def get_paginated(client: Client, url: str, **kwargs: str | int) -> Iterator[lis
2827
the full result set is returned.
2928
3029
"""
31-
3230
while True:
3331
response = client.get(url, params={**kwargs})
3432
response.raise_for_status()

rdw_ingestion_tools/api/flow_results/main.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44
from api.flow_results.requests.flows import Flows
55
from api.flow_results.requests.responses import Responses
66

7-
from . import client as default_client
7+
from .client import get_client
88

99

1010
@define
@@ -15,7 +15,7 @@ class pyFlows:
1515
1616
"""
1717

18-
client: Client = field(factory=lambda: default_client)
18+
client: Client = field(factory=get_client)
1919

2020
flows: Flows = field(init=False)
2121
responses: Responses = field(init=False)
Lines changed: 3 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -1,17 +1,4 @@
1-
from httpx import Client
2-
3-
from .. import config_from_env
4-
5-
API_KEY = config_from_env("TURN_BQ_API_KEY")
6-
BASE_URL = config_from_env("TURN_BQ_API_BASE_URL")
7-
8-
session_headers = {
9-
"Authorization": f"Bearer {API_KEY}",
10-
"Accept": "application/vnd.v1+json",
11-
"Content-Type": "application/json",
12-
}
13-
14-
client: Client = Client(base_url=BASE_URL, headers=session_headers, timeout=30.0)
15-
16-
1+
from .client import get_client, make_client
172
from .main import pyTurnBQ as pyTurnBQ
3+
4+
__all__ = ["get_client", "make_client", "pyTurnBQ"]
Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
from functools import lru_cache
2+
3+
from httpx import Client
4+
from httpx_retries import RetryTransport
5+
6+
from .. import config_from_env
7+
8+
API_KEY = config_from_env("TURN_BQ_API_KEY")
9+
BASE_URL = config_from_env("TURN_BQ_API_BASE_URL")
10+
11+
session_headers = {
12+
"Authorization": f"Bearer {API_KEY}",
13+
"Accept": "application/vnd.v1+json",
14+
"Content-Type": "application/json",
15+
}
16+
17+
18+
def make_client() -> Client:
19+
return Client(
20+
base_url=BASE_URL,
21+
headers=session_headers,
22+
timeout=30.0,
23+
transport=RetryTransport(),
24+
)
25+
26+
27+
@lru_cache(maxsize=1)
28+
def get_client() -> Client:
29+
return make_client()

rdw_ingestion_tools/api/turn_bq/main.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@
1111
from api.turn_bq.requests.messages import Messages
1212
from api.turn_bq.requests.statuses import Statuses
1313

14-
from . import client as default_client
14+
from .client import get_client
1515

1616

1717
@define
@@ -22,7 +22,7 @@ class pyTurnBQ:
2222
2323
"""
2424

25-
client: Client = field(factory=lambda: default_client)
25+
client: Client = field(factory=get_client)
2626

2727
cards: Cards = field(init=False)
2828
contacts: Contacts = field(init=False)
Lines changed: 206 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,206 @@
1+
import importlib
2+
from unittest.mock import MagicMock, patch
3+
4+
import httpx
5+
import pytest
6+
7+
8+
@pytest.fixture(autouse=True)
9+
def _set_required_env(monkeypatch):
10+
"""Ensure required env vars exist to avoid import-time errors."""
11+
monkeypatch.setenv("FLOW_RESULTS_API_KEY", "test-key")
12+
monkeypatch.setenv("FLOW_RESULTS_API_BASE_URL", "http://test-api.com")
13+
14+
15+
@pytest.fixture
16+
def flow_results_httpx_module(_set_required_env):
17+
"""Import the module under test."""
18+
return importlib.import_module(
19+
"rdw_ingestion_tools.api.flow_results.extensions.httpx"
20+
)
21+
22+
23+
@pytest.fixture
24+
def mock_client(monkeypatch, flow_results_httpx_module):
25+
"""
26+
Patch the Client symbol used in the module under test
27+
and return the mock client.
28+
"""
29+
client = MagicMock()
30+
client.__enter__.return_value = client
31+
client.__exit__.return_value = False
32+
33+
monkeypatch.setattr(
34+
flow_results_httpx_module, "Client", lambda *args, **kwargs: client
35+
)
36+
return client
37+
38+
39+
@pytest.fixture
40+
def make_response():
41+
"""
42+
Factory to construct a mock Response with a provided JSON payload
43+
and optional status error.
44+
"""
45+
46+
def _factory(json_data: dict, status_error: Exception | None = None):
47+
resp = MagicMock()
48+
resp.json.return_value = json_data
49+
if status_error is None:
50+
resp.raise_for_status.return_value = None
51+
else:
52+
resp.raise_for_status.side_effect = status_error
53+
return resp
54+
55+
return _factory
56+
57+
58+
def test_get_ids(flow_results_httpx_module, mock_client, make_response):
59+
json_payload = {
60+
"data": [
61+
{"id": "f-1"},
62+
{"id": "f-2"},
63+
]
64+
}
65+
mock_client.get.return_value = make_response(json_payload)
66+
67+
result = list(flow_results_httpx_module.get_ids(mock_client, org="acme", page=1))
68+
69+
assert result == ["f-1", "f-2"]
70+
mock_client.get.assert_called_once_with("", params={"org": "acme", "page": 1})
71+
72+
73+
def test_get_ids_error(flow_results_httpx_module, mock_client, make_response):
74+
req = httpx.Request("GET", "http://example.invalid")
75+
resp = httpx.Response(500, request=req)
76+
status_err = httpx.HTTPStatusError("Error", request=req, response=resp)
77+
mock_client.get.return_value = make_response({"data": []}, status_error=status_err)
78+
79+
with pytest.raises(httpx.HTTPStatusError):
80+
list(flow_results_httpx_module.get_ids(mock_client))
81+
82+
83+
def test_get_paginated_single_page(
84+
flow_results_httpx_module, mock_client, make_response
85+
):
86+
json_payload = {
87+
"data": {
88+
"attributes": {"responses": [[{"a": 1}], [{"b": 2}]]},
89+
"relationships": {"links": {"next": None}},
90+
}
91+
}
92+
# Accessing ["next"] on a dict with None value won't raise AttributeError,
93+
# but our code only checks AttributeError. To keep a single page, we'll instead
94+
# make the relationships object raise AttributeError on __getitem__
95+
# in the pagination test.
96+
# For single-page, just ensure we don't provide a usable "next" URL
97+
# and call count is 1.
98+
mock_client.get.return_value = make_response(json_payload)
99+
100+
result = list(flow_results_httpx_module.get_paginated(mock_client, "/packages/123"))
101+
102+
assert result == [[{"a": 1}], [{"b": 2}]]
103+
mock_client.get.assert_called_once_with("/packages/123", params={})
104+
105+
106+
def test_get_paginated_pagination(
107+
flow_results_httpx_module, mock_client, make_response
108+
):
109+
# First page returns a next URL
110+
first_page = {
111+
"data": {
112+
"attributes": {"responses": [[{"a": 1}]]},
113+
"relationships": {
114+
"links": {"next": "https://api.example.com/packages/next-token"}
115+
},
116+
}
117+
}
118+
119+
# Second page should trigger the break via AttributeError when accessing ["next"]
120+
bad_links = MagicMock()
121+
bad_links.__getitem__.side_effect = AttributeError
122+
second_page = {
123+
"data": {
124+
"attributes": {"responses": [[{"c": 3}]]},
125+
"relationships": bad_links,
126+
}
127+
}
128+
129+
mock_client.get.side_effect = [
130+
make_response(first_page),
131+
make_response(second_page),
132+
]
133+
134+
result = list(
135+
flow_results_httpx_module.get_paginated(mock_client, "/packages/start", q="x")
136+
)
137+
138+
assert result == [[{"a": 1}], [{"c": 3}]]
139+
assert mock_client.get.call_count == 2
140+
mock_client.get.assert_any_call("/packages/start", params={"q": "x"})
141+
# The second call should use the path part after "packages/"
142+
mock_client.get.assert_any_call("next-token", params={"q": "x"})
143+
144+
145+
def test_get_paginated_kwargs_propagation(
146+
flow_results_httpx_module, mock_client, make_response
147+
):
148+
json_payload = {
149+
"data": {
150+
"attributes": {"responses": []},
151+
"relationships": {"links": {"next": None}},
152+
}
153+
}
154+
mock_client.get.return_value = make_response(json_payload)
155+
156+
list(
157+
flow_results_httpx_module.get_paginated(
158+
mock_client, "/packages/xyz", limit=50, cursor="abc"
159+
)
160+
)
161+
162+
mock_client.get.assert_called_once_with(
163+
"/packages/xyz", params={"limit": 50, "cursor": "abc"}
164+
)
165+
166+
167+
def test_get_paginated_retry_mechanism(
168+
flow_results_httpx_module, mock_client, make_response
169+
):
170+
"""Test that get_paginated uses RetryTransport for resilient HTTP requests."""
171+
# Create a mock response
172+
json_payload = {
173+
"data": {
174+
"attributes": {"responses": [[{"a": 1}], [{"b": 2}]]},
175+
"relationships": {"links": {"next": None}},
176+
}
177+
}
178+
mock_response = make_response(json_payload)
179+
mock_client.get.return_value = mock_response
180+
181+
# Mock the RetryTransport using the correct import path
182+
with patch(
183+
"rdw_ingestion_tools.api.flow_results.client.Client",
184+
return_value=mock_client,
185+
):
186+
# Call the function
187+
result = list(
188+
flow_results_httpx_module.get_paginated(
189+
mock_client, "http://test-api.com/data"
190+
)
191+
)
192+
193+
# Assertions
194+
assert len(result) == 2
195+
assert result == [[{"a": 1}], [{"b": 2}]]
196+
197+
# Verify the call was made with correct parameters
198+
mock_client.get.assert_called_once_with("http://test-api.com/data", params={})
199+
200+
201+
def test_make_client_uses_retrytransport():
202+
with patch("rdw_ingestion_tools.api.flow_results.client.RetryTransport") as mock_rt:
203+
from rdw_ingestion_tools.api.flow_results.client import make_client
204+
205+
_ = make_client()
206+
assert mock_rt.call_count == 1

0 commit comments

Comments
 (0)