|
4 | 4 |
|
5 | 5 | import json
|
6 | 6 |
|
| 7 | +import pytest |
7 | 8 | import requests
|
8 | 9 | from airbyte_cdk.sources.declarative.decoders.json_decoder import JsonDecoder
|
9 | 10 | from airbyte_cdk.sources.declarative.extractors.jello import JelloExtractor
|
10 | 11 |
|
11 | 12 | config = {"field": "record_array"}
|
12 |
| -decoder = JsonDecoder() |
13 |
| - |
14 |
| - |
15 |
| -def test(): |
16 |
| - transform = "_.data" |
17 |
| - extractor = JelloExtractor(transform, decoder, config) |
18 |
| - |
19 |
| - records = [{"id": 1}, {"id": 2}] |
20 |
| - body = {"data": records} |
21 |
| - response = create_response(body) |
22 |
| - actual_records = extractor.extract_records(response) |
23 |
| - |
24 |
| - assert actual_records == records |
25 |
| - |
26 |
| - |
27 |
| -def test_field_in_config(): |
28 |
| - transform = "_.{{ config['field'] }}" |
29 |
| - extractor = JelloExtractor(transform, decoder, config) |
| 13 | +kwargs = {"data_field": "records"} |
30 | 14 |
|
31 |
| - records = [{"id": 1}, {"id": 2}] |
32 |
| - body = {"record_array": records} |
33 |
| - response = create_response(body) |
34 |
| - actual_records = extractor.extract_records(response) |
35 |
| - |
36 |
| - assert actual_records == records |
| 15 | +decoder = JsonDecoder() |
37 | 16 |
|
38 | 17 |
|
39 |
| -def test_field_in_kwargs(): |
40 |
| - transform = "_.{{ kwargs['data_field'] }}" |
41 |
| - kwargs = {"data_field": "records"} |
| 18 | +@pytest.mark.parametrize( |
| 19 | + "test_name, transform, body, expected_records", |
| 20 | + [ |
| 21 | + ("test_extract_from_array", "_.data", {"data": [{"id": 1}, {"id": 2}]}, [{"id": 1}, {"id": 2}]), |
| 22 | + ("test_field_in_config", "_.{{ config['field'] }}", {"record_array": [{"id": 1}, {"id": 2}]}, [{"id": 1}, {"id": 2}]), |
| 23 | + ("test_field_in_kwargs", "_.{{ kwargs['data_field'] }}", {"records": [{"id": 1}, {"id": 2}]}, [{"id": 1}, {"id": 2}]), |
| 24 | + ("test_default", "_{{kwargs['field']}}", [{"id": 1}, {"id": 2}], [{"id": 1}, {"id": 2}]), |
| 25 | + ( |
| 26 | + "test_remove_fields_from_records", |
| 27 | + "[{k:v for k,v in d.items() if k != 'value_to_remove'} for d in _.data]", |
| 28 | + {"data": [{"id": 1, "value": "HELLO", "value_to_remove": "fail"}, {"id": 2, "value": "WORLD", "value_to_remove": "fail"}]}, |
| 29 | + [{"id": 1, "value": "HELLO"}, {"id": 2, "value": "WORLD"}], |
| 30 | + ), |
| 31 | + ( |
| 32 | + "test_add_fields_from_records", |
| 33 | + "[{**{k:v for k,v in d.items()}, **{'project_id': d['project']['id']}} for d in _.data]", |
| 34 | + {"data": [{"id": 1, "value": "HELLO", "project": {"id": 8}}, {"id": 2, "value": "WORLD", "project": {"id": 9}}]}, |
| 35 | + [ |
| 36 | + {"id": 1, "value": "HELLO", "project_id": 8, "project": {"id": 8}}, |
| 37 | + {"id": 2, "value": "WORLD", "project_id": 9, "project": {"id": 9}}, |
| 38 | + ], |
| 39 | + ), |
| 40 | + ], |
| 41 | +) |
| 42 | +def test(test_name, transform, body, expected_records): |
42 | 43 | extractor = JelloExtractor(transform, decoder, config, kwargs=kwargs)
|
43 | 44 |
|
44 |
| - records = [{"id": 1}, {"id": 2}] |
45 |
| - body = {"records": records} |
46 | 45 | response = create_response(body)
|
47 | 46 | actual_records = extractor.extract_records(response)
|
48 | 47 |
|
49 |
| - assert actual_records == records |
| 48 | + assert actual_records == expected_records |
50 | 49 |
|
51 | 50 |
|
52 | 51 | def create_response(body):
|
53 | 52 | response = requests.Response()
|
54 | 53 | response._content = json.dumps(body).encode("utf-8")
|
55 | 54 | return response
|
56 |
| - |
57 |
| - |
58 |
| -def test_default(): |
59 |
| - transform = "_{{kwargs['field']}}" |
60 |
| - extractor = JelloExtractor(transform, decoder, config) |
61 |
| - |
62 |
| - records = [{"id": 1}, {"id": 2}] |
63 |
| - response = create_response(records) |
64 |
| - actual_records = extractor.extract_records(response) |
65 |
| - |
66 |
| - assert actual_records == records |
0 commit comments