-
Notifications
You must be signed in to change notification settings - Fork 1.3k
Expand file tree
/
Copy pathtest_arrow_error_decorator.py
More file actions
231 lines (185 loc) · 8.22 KB
/
test_arrow_error_decorator.py
File metadata and controls
231 lines (185 loc) · 8.22 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
from unittest.mock import MagicMock, patch
import pyarrow.flight as fl
import pytest
from pydantic import ValidationError
from feast.arrow_error_handler import (
_get_exception_data,
arrow_client_error_handling_decorator,
)
from feast.errors import FeatureViewNotFoundException, PermissionNotFoundException
from feast.infra.offline_stores.remote import RemoteOfflineStoreConfig
permissionError = PermissionNotFoundException("dummy_name", "dummy_project")
@arrow_client_error_handling_decorator
def decorated_method(error):
raise error
@pytest.mark.parametrize(
"error, expected_raised_error",
[
(fl.FlightError("Flight error: "), fl.FlightError("Flight error: ")),
(
fl.FlightError(f"Flight error: {permissionError.to_error_detail()}"),
permissionError,
),
(fl.FlightError("Test Error"), fl.FlightError("Test Error")),
(RuntimeError("Flight error: "), RuntimeError("Flight error: ")),
(permissionError, permissionError),
],
)
def test_rest_error_handling_with_feast_exception(error, expected_raised_error):
with pytest.raises(
type(expected_raised_error),
match=str(expected_raised_error),
):
decorated_method(error)
class TestArrowClientRetry:
@patch("feast.arrow_error_handler.time.sleep")
def test_retries_on_flight_unavailable_error(self, mock_sleep):
client = MagicMock()
client._connection_retries = 3
call_count = 0
@arrow_client_error_handling_decorator
def flaky_method(self_arg):
nonlocal call_count
call_count += 1
if call_count < 3:
raise fl.FlightUnavailableError("Connection refused")
return "success"
result = flaky_method(client)
assert result == "success"
assert call_count == 3
assert mock_sleep.call_count == 2
@patch("feast.arrow_error_handler.time.sleep")
def test_raises_after_max_retries_exhausted(self, mock_sleep):
client = MagicMock()
client._connection_retries = 3
@arrow_client_error_handling_decorator
def always_unavailable(self_arg):
raise fl.FlightUnavailableError("Connection refused")
with pytest.raises(fl.FlightUnavailableError, match="Connection refused"):
always_unavailable(client)
assert mock_sleep.call_count == 3
@patch("feast.arrow_error_handler.time.sleep")
def test_respects_connection_retries_from_client(self, mock_sleep):
client = MagicMock()
client._connection_retries = 1
call_count = 0
@arrow_client_error_handling_decorator
def method_on_client(self_arg):
nonlocal call_count
call_count += 1
raise fl.FlightUnavailableError("Connection refused")
with pytest.raises(fl.FlightUnavailableError):
method_on_client(client)
assert call_count == 2 # 1 initial + 1 retry
assert mock_sleep.call_count == 1
@patch("feast.arrow_error_handler.time.sleep")
def test_no_retry_on_non_transient_errors(self, mock_sleep):
client = MagicMock()
client._connection_retries = 3
call_count = 0
@arrow_client_error_handling_decorator
def method_with_error(self_arg):
nonlocal call_count
call_count += 1
raise fl.FlightError("Permanent error")
with pytest.raises(fl.FlightError, match="Permanent error"):
method_with_error(client)
assert call_count == 1
mock_sleep.assert_not_called()
@patch("feast.arrow_error_handler.time.sleep")
def test_exponential_backoff_timing(self, mock_sleep):
client = MagicMock()
client._connection_retries = 3
@arrow_client_error_handling_decorator
def always_unavailable(self_arg):
raise fl.FlightUnavailableError("Connection refused")
with pytest.raises(fl.FlightUnavailableError):
always_unavailable(client)
wait_times = [call.args[0] for call in mock_sleep.call_args_list]
assert wait_times == [0.5, 1.0, 2.0]
@patch("feast.arrow_error_handler.time.sleep")
def test_zero_retries_disables_retry(self, mock_sleep):
client = MagicMock()
client._connection_retries = 0
call_count = 0
@arrow_client_error_handling_decorator
def method_on_client(self_arg):
nonlocal call_count
call_count += 1
raise fl.FlightUnavailableError("Connection refused")
with pytest.raises(fl.FlightUnavailableError):
method_on_client(client)
assert call_count == 1
mock_sleep.assert_not_called()
@patch("feast.arrow_error_handler.time.sleep")
def test_no_retry_for_standalone_stream_functions(self, mock_sleep):
"""Standalone functions (write_table, read_all) where args[0] is a
writer/reader should not retry since broken streams can't be reused."""
writer = MagicMock(spec=[]) # no _connection_retries attribute
call_count = 0
@arrow_client_error_handling_decorator
def write_table(w):
nonlocal call_count
call_count += 1
raise fl.FlightUnavailableError("stream broken")
with pytest.raises(fl.FlightUnavailableError, match="stream broken"):
write_table(writer)
assert call_count == 1
mock_sleep.assert_not_called()
@patch("feast.arrow_error_handler.time.sleep")
def test_negative_connection_retries_treated_as_zero(self, mock_sleep):
"""Negative _connection_retries must not skip function execution."""
client = MagicMock()
client._connection_retries = -1
call_count = 0
@arrow_client_error_handling_decorator
def method_on_client(self_arg):
nonlocal call_count
call_count += 1
return "ok"
result = method_on_client(client)
assert result == "ok"
assert call_count == 1
mock_sleep.assert_not_called()
def test_config_rejects_negative_connection_retries(self):
with pytest.raises(ValidationError):
RemoteOfflineStoreConfig(host="localhost", connection_retries=-1)
class TestGetExceptionData:
def test_non_string_input_returns_empty(self):
assert _get_exception_data(12345) == ""
assert _get_exception_data(None) == ""
assert _get_exception_data(b"bytes") == ""
def test_no_flight_error_prefix_returns_empty(self):
assert _get_exception_data("some random error") == ""
def test_flight_error_prefix_without_json_returns_empty(self):
assert _get_exception_data("Flight error: no json here") == ""
def test_extracts_json_from_flight_error(self):
fv_error = FeatureViewNotFoundException("my_view", "my_project")
error_str = f"Flight error: {fv_error.to_error_detail()}"
result = _get_exception_data(error_str)
assert '"class": "FeatureViewNotFoundException"' in result
assert '"module": "feast.errors"' in result
def test_extracts_json_with_trailing_text(self):
fv_error = FeatureViewNotFoundException("my_view", "my_project")
error_str = (
f"Flight error: {fv_error.to_error_detail()}. "
"gRPC client debug context: some extra info"
)
result = _get_exception_data(error_str)
assert '"class": "FeatureViewNotFoundException"' in result
assert '"module": "feast.errors"' in result
def test_extracts_json_with_grpc_debug_context_containing_braces(self):
fv_error = FeatureViewNotFoundException("my_view", "my_project")
error_str = (
f"Flight error: {fv_error.to_error_detail()}. "
"gRPC client debug context: UNKNOWN:Error received from peer "
'ipv4:127.0.0.1:59930 {grpc_message:"Flight error: ...", '
'grpc_status:2, created_time:"2026-03-17T17:32:07"}'
)
result = _get_exception_data(error_str)
assert '"class": "FeatureViewNotFoundException"' in result
assert '"module": "feast.errors"' in result
from feast.errors import FeastError
reconstructed = FeastError.from_error_detail(result)
assert reconstructed is not None
assert "my_view" in str(reconstructed)