-
Notifications
You must be signed in to change notification settings - Fork 109
Expand file tree
/
Copy pathhelper.py
More file actions
397 lines (312 loc) · 11.7 KB
/
helper.py
File metadata and controls
397 lines (312 loc) · 11.7 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
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
"""
Helper utilities for platform tests.
Provides:
- Lightweight REST wrappers (no SDK abstraction where raw status codes matter)
- Polling helpers (compilation, generic condition)
- Simple selector/object helpers
No automatic cleanup; pipelines are left in place for inspection after failures.
"""
from __future__ import annotations
import time
import json
import logging
import pytest
import requests
from dataclasses import dataclass
from typing import Any, Dict, Iterable, Optional
from http import HTTPStatus
from urllib.parse import quote, quote_plus
from tests import FELDERA_TLS_INSECURE, API_KEY, BASE_URL, unique_pipeline_name
from feldera.testutils_oidc import get_oidc_test_helper
API_PREFIX = "/v0"
def _base_headers() -> Dict[str, str]:
headers = {
"Accept": "application/json",
}
# Try OIDC authentication first, then fall back to API_KEY
oidc_helper = get_oidc_test_helper()
if oidc_helper is not None:
token = oidc_helper.obtain_access_token()
headers["Authorization"] = f"Bearer {token}"
elif API_KEY:
headers["Authorization"] = f"Bearer {API_KEY}"
return headers
def api_url(fragment: str) -> str:
if not fragment.startswith("/"):
fragment = "/" + fragment
return f"{API_PREFIX}{fragment}"
def http_request(method: str, path: str, **kwargs) -> requests.Response:
"""
Low-level request wrapper (no retries). Raises only on network errors.
Args:
method: HTTP method (GET, POST, etc.)
path: URL path
base_headers: Optional override for base headers. If None (default), uses _base_headers().
If empty dict {}, no base headers are applied (for testing unauthenticated requests).
**kwargs: Additional arguments passed to requests.request()
"""
if not path.startswith("/"):
path = "/" + path
url = BASE_URL.rstrip("/") + path
# Allow override of base headers for testing unauthenticated requests
base_headers_arg = kwargs.pop("base_headers", None)
if base_headers_arg is None:
base_headers = _base_headers() # Default: include auth headers
else:
base_headers = base_headers_arg # Override: could be {} for no auth
custom_headers = kwargs.pop("headers", None) or {}
headers = {
**base_headers,
**custom_headers,
} # Merge, with custom headers taking precedence
# Provide a default timeout to avoid hanging tests.
timeout = kwargs.pop("timeout", 30)
kwargs["verify"] = not FELDERA_TLS_INSECURE
resp = requests.request(
method.upper(), url, headers=headers, timeout=timeout, **kwargs
)
return resp
def get(path: str, **kw) -> requests.Response:
return http_request("GET", path, **kw)
def get_pipeline(name: str, selector: str) -> requests.Response:
return get(f"{API_PREFIX}/pipelines/{name}?selector={selector}")
def post_no_body(path: str, **kw):
return http_request("POST", path, **kw)
def post_json(path: str, body: Dict[str, Any], **kw) -> requests.Response:
return http_request("POST", path, json=body, **kw)
def put_json(path: str, body: Dict[str, Any], **kw) -> requests.Response:
return http_request("PUT", path, json=body, **kw)
def patch_json(path: str, body: Dict[str, Any], **kw) -> requests.Response:
return http_request("PATCH", path, json=body, **kw)
def delete(path: str, **kw) -> requests.Response:
return http_request("DELETE", path, **kw)
def wait_for_deployment_status(name: str, desired: str, timeout_s: float = 60.0):
deadline = time.time() + timeout_s
last = None
while time.time() < deadline:
r = get_pipeline(name, "status")
if r.status_code != HTTPStatus.OK:
time.sleep(0.25)
continue
obj = r.json()
last = obj.get("deployment_status")
if last == desired:
return obj
time.sleep(0.25)
raise TimeoutError(
f"Timed out waiting for pipeline '{name}' deployment_status={desired} (last={last})"
)
def start_pipeline(name: str, wait: bool = True):
r = post_no_body(api_url(f"/pipelines/{name}/start"))
assert r.status_code == HTTPStatus.ACCEPTED, (
f"Unexpected start response: {r.status_code} {r.text}"
)
if wait:
wait_for_deployment_status(name, "Running", 30)
return r
def resume_pipeline(name: str, wait: bool = True):
r = post_no_body(api_url(f"/pipelines/{name}/resume"))
assert r.status_code == HTTPStatus.ACCEPTED, (
f"Unexpected resume response: {r.status_code} {r.text}"
)
if wait:
wait_for_deployment_status(name, "Running", 30)
return r
def start_pipeline_as_paused(name: str, wait: bool = True):
r = post_no_body(api_url(f"/pipelines/{name}/start"), params={"initial": "paused"})
assert r.status_code == HTTPStatus.ACCEPTED, (
f"Unexpected pause response: {r.status_code} {r.text}"
)
if wait:
wait_for_deployment_status(name, "Paused", 30)
return r
def pause_pipeline(name: str, wait: bool = True):
r = post_no_body(api_url(f"/pipelines/{name}/pause"))
assert r.status_code == HTTPStatus.ACCEPTED, (
f"Unexpected pause response: {r.status_code} {r.text}"
)
if wait:
wait_for_deployment_status(name, "Paused", 30)
return r
def stop_pipeline(name: str, force: bool = True, wait: bool = True):
r = post_no_body(
api_url(f"/pipelines/{name}/stop?force={'true' if force else 'false'}")
)
assert r.status_code == HTTPStatus.ACCEPTED, (
f"Unexpected stop response: {r.status_code} {r.text}"
)
wait_for_deployment_status(name, "Stopped", 30)
return r
def wait_for_cleared_storage(name: str, timeout_s: float = 60.0):
deadline = time.time() + timeout_s
last = None
while time.time() < deadline:
r = get_pipeline(name, "status")
if r.status_code != HTTPStatus.OK:
time.sleep(0.25)
continue
obj = r.json()
last = obj.get("storage_status")
if last == "Cleared":
return obj
time.sleep(0.25)
raise TimeoutError(
f"Timed out waiting for pipeline '{name}' to clear storage (last={last})"
)
def clear_pipeline(name: str, wait: bool = True):
r = post_no_body(f"{API_PREFIX}/pipelines/{name}/clear")
if wait and r.status_code == HTTPStatus.ACCEPTED:
wait_for_cleared_storage(name)
return r
def reset_pipeline(name: str):
if get_pipeline(name, "status").status_code != HTTPStatus.OK:
return
stop_pipeline(name, force=True)
clear_pipeline(name)
def delete_pipeline(name: str):
return delete(api_url(f"pipelines/{name}"))
def cleanup_pipeline(name: str):
reset_pipeline(name)
delete_pipeline(name)
def connector_action(pipeline: str, table: str, connector: str, action: str):
table_enc = quote(table, safe="")
r = post_no_body(
api_url(
f"/pipelines/{pipeline}/tables/{table_enc}/connectors/{connector}/{action}"
)
)
return r
def adhoc_query_json(pipeline: str, sql: str):
"Runs a SQL query, returns results as list of dicts (JSON lines format)."
path = api_url(f"/pipelines/{pipeline}/query?sql={quote_plus(sql)}&format=json")
resp = get(path)
assert resp.status_code == HTTPStatus.OK, (
f"Adhoc query failed: status={resp.status_code}, body={resp.text}"
)
raw = resp.text.strip()
if not raw:
return []
lines = raw.split("\n")
return [json.loads(line) for line in lines if line]
def pipeline_stats(pipeline: str):
r = get(api_url(f"/pipelines/{pipeline}/stats"))
assert r.status_code == HTTPStatus.OK, (r.status_code, r.text)
return r.json()
def connector_stats(pipeline: str, table: str, connector: str):
table_enc = quote(table, safe="")
res = get(
api_url(
f"/pipelines/{pipeline}/tables/{table_enc}/connectors/{connector}/stats"
)
)
assert res.status_code == HTTPStatus.OK, (res.status_code, res.text)
return res.json()
def connector_paused(pipeline, table: str, connector: str) -> bool:
return connector_stats(pipeline, table, connector)["paused"]
@dataclass
class WaitResult:
ok: bool
last_status: Optional[str]
last_object: Optional[Dict[str, Any]]
elapsed_s: float
def wait_for_program_success(
pipeline_name: str,
expected_program_version: int,
timeout_s: float = 1800.0,
sleep_s: float = 0.5,
) -> WaitResult:
"""
Poll until the pipeline's program_status is Success and program_version
>= expected_program_version.
Mirrors semantics of the Rust `wait_for_compiled_program` helper.
Returns a WaitResult. Raises AssertionError on compile error, TimeoutError on timeout.
"""
deadline = time.time() + timeout_s
last_status = None
last_obj: Optional[Dict[str, Any]] = None
while True:
if time.time() > deadline:
raise TimeoutError(
f"Timed out waiting for pipeline '{pipeline_name}' to compile "
f"(expected program_version >= {expected_program_version}, "
f"last_status={last_status}, last_obj={last_obj})"
)
resp = get(f"{API_PREFIX}/pipelines/{pipeline_name}")
if resp.status_code == HTTPStatus.NOT_FOUND:
raise RuntimeError(
f"Pipeline '{pipeline_name}' disappeared during compilation wait"
)
try:
obj = resp.json()
except Exception as e:
raise RuntimeError(
f"Failed to parse pipeline JSON: {e}; body={resp.text!r}"
)
last_obj = obj
last_status = obj.get("program_status")
version = obj.get("program_version") or 0
if last_status == "Success" and version >= expected_program_version:
return WaitResult(
True, last_status, last_obj, timeout_s - (deadline - time.time())
)
if last_status in ("SqlError", "RustError"):
raise AssertionError(
"Compilation failed: " + json.dumps(obj.get("program_error"), indent=2)
)
time.sleep(sleep_s)
def wait_for_condition(
description: str,
predicate,
timeout_s: float = 30.0,
sleep_s: float = 0.2,
) -> None:
"""
Generic polling helper.
predicate: callable returning truthy when condition met (can be sync or async).
"""
start = time.time()
deadline = start + timeout_s
attempt = 0
while True:
now = time.time()
if now > deadline:
raise TimeoutError(f"Timeout waiting for condition: {description}")
attempt += 1
try:
result = predicate()
except Exception as e: # noqa: BLE001
logging.debug("Predicate raised %s (attempt %d), continuing", e, attempt)
result = False
if result:
return
time.sleep(sleep_s)
def extract_object_by_name(
collection: Iterable[Dict[str, Any]], name: str
) -> Dict[str, Any]:
for obj in collection:
if obj.get("name") == name:
return obj
raise KeyError(f"Object with name '{name}' not found in collection")
def gen_pipeline_name(func):
"""
Decorator for pytest functions that automatically generates a unique pipeline name.
The decorated function will receive a 'pipeline_name' parameter.
After the test completes, attempts to delete the pipeline but ignores any errors.
"""
return pytest.mark.usefixtures("pipeline_name")(func)
__all__ = [
"API_PREFIX",
"HTTPStatus",
"http_request",
"get",
"post_json",
"put_json",
"patch_json",
"delete",
"wait_for_program_success",
"wait_for_condition",
"unique_pipeline_name",
"extract_object_by_name",
"gen_pipeline_name",
"connector_action",
]