Skip to content

Commit 4f7a84e

Browse files
[cherry-pick-me]: use kw args for Test* dataclasses
1 parent cca55f7 commit 4f7a84e

File tree

6 files changed

+115
-42
lines changed

6 files changed

+115
-42
lines changed

airbyte_cdk/connector_builder/connector_builder_handler.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -86,7 +86,9 @@ def read_stream(
8686
) -> AirbyteMessage:
8787
try:
8888
test_read_handler = TestReader(
89-
limits.max_pages_per_slice, limits.max_slices, limits.max_records
89+
max_pages_per_slice=limits.max_pages_per_slice,
90+
max_slices=limits.max_slices,
91+
max_record_limit=limits.max_records,
9092
)
9193
# The connector builder only supports a single stream
9294
stream_name = configured_catalog.streams[0].stream.name

unit_tests/connector_builder/test_connector_builder_handler.py

Lines changed: 20 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1069,7 +1069,11 @@ def test_read_source(mock_http_stream):
10691069
max_records = 100
10701070
max_pages_per_slice = 2
10711071
max_slices = 3
1072-
limits = TestLimits(max_records, max_pages_per_slice, max_slices)
1072+
limits = TestLimits(
1073+
max_records=max_records,
1074+
max_pages_per_slice=max_pages_per_slice,
1075+
max_slices=max_slices,
1076+
)
10731077

10741078
catalog = ConfiguredAirbyteCatalog(
10751079
streams=[
@@ -1116,7 +1120,11 @@ def test_read_source_single_page_single_slice(mock_http_stream):
11161120
max_records = 100
11171121
max_pages_per_slice = 1
11181122
max_slices = 1
1119-
limits = TestLimits(max_records, max_pages_per_slice, max_slices)
1123+
limits = TestLimits(
1124+
max_records=max_records,
1125+
max_pages_per_slice=max_pages_per_slice,
1126+
max_slices=max_slices,
1127+
)
11201128

11211129
catalog = ConfiguredAirbyteCatalog(
11221130
streams=[
@@ -1200,7 +1208,11 @@ def test_handle_read_external_requests(deployment_mode, url_base, expected_error
12001208
endpoints when running on Cloud or OSS deployments
12011209
"""
12021210

1203-
limits = TestLimits(max_records=100, max_pages_per_slice=1, max_slices=1)
1211+
limits = TestLimits(
1212+
max_records=100,
1213+
max_pages_per_slice=1,
1214+
max_slices=1,
1215+
)
12041216

12051217
catalog = ConfiguredAirbyteCatalog(
12061218
streams=[
@@ -1286,7 +1298,11 @@ def test_handle_read_external_oauth_request(deployment_mode, token_url, expected
12861298
endpoints when running on Cloud or OSS deployments
12871299
"""
12881300

1289-
limits = TestLimits(max_records=100, max_pages_per_slice=1, max_slices=1)
1301+
limits = TestLimits(
1302+
max_records=100,
1303+
max_pages_per_slice=1,
1304+
max_slices=1,
1305+
)
12901306

12911307
catalog = ConfiguredAirbyteCatalog(
12921308
streams=[

unit_tests/connector_builder/test_message_grouper.py

Lines changed: 70 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -203,7 +203,10 @@ def test_get_grouped_messages(mock_entrypoint_read: Mock) -> None:
203203
),
204204
)
205205

206-
connector_builder_handler = TestReader(MAX_PAGES_PER_SLICE, MAX_SLICES)
206+
connector_builder_handler = TestReader(
207+
max_pages_per_slice=MAX_PAGES_PER_SLICE,
208+
max_slices=MAX_SLICES,
209+
)
207210
actual_response: StreamRead = connector_builder_handler.run_test_read(
208211
source=mock_source,
209212
config=CONFIG,
@@ -287,7 +290,10 @@ def test_get_grouped_messages_with_logs(mock_entrypoint_read: Mock) -> None:
287290
),
288291
)
289292

290-
connector_builder_handler = TestReader(MAX_PAGES_PER_SLICE, MAX_SLICES)
293+
connector_builder_handler = TestReader(
294+
max_pages_per_slice=MAX_PAGES_PER_SLICE,
295+
max_slices=MAX_SLICES,
296+
)
291297

292298
actual_response: StreamRead = connector_builder_handler.run_test_read(
293299
source=mock_source,
@@ -340,7 +346,11 @@ def test_get_grouped_messages_record_limit(
340346
n_records = 2
341347
record_limit = min(request_record_limit, max_record_limit)
342348

343-
api = TestReader(MAX_PAGES_PER_SLICE, MAX_SLICES, max_record_limit=max_record_limit)
349+
api = TestReader(
350+
max_pages_per_slice=MAX_PAGES_PER_SLICE,
351+
max_slices=MAX_SLICES,
352+
max_record_limit=max_record_limit,
353+
)
344354
# this is the call we expect to raise an exception
345355
if should_fail:
346356
with pytest.raises(ValueError):
@@ -404,7 +414,11 @@ def test_get_grouped_messages_default_record_limit(
404414
)
405415
n_records = 2
406416

407-
api = TestReader(MAX_PAGES_PER_SLICE, MAX_SLICES, max_record_limit=max_record_limit)
417+
api = TestReader(
418+
max_pages_per_slice=MAX_PAGES_PER_SLICE,
419+
max_slices=MAX_SLICES,
420+
max_record_limit=max_record_limit,
421+
)
408422
actual_response: StreamRead = api.run_test_read(
409423
source=mock_source,
410424
config=CONFIG,
@@ -443,7 +457,10 @@ def test_get_grouped_messages_limit_0(mock_entrypoint_read: Mock) -> None:
443457
]
444458
),
445459
)
446-
api = TestReader(MAX_PAGES_PER_SLICE, MAX_SLICES)
460+
api = TestReader(
461+
max_pages_per_slice=MAX_PAGES_PER_SLICE,
462+
max_slices=MAX_SLICES,
463+
)
447464

448465
with pytest.raises(ValueError):
449466
api.run_test_read(
@@ -501,7 +518,10 @@ def test_get_grouped_messages_no_records(mock_entrypoint_read: Mock) -> None:
501518
),
502519
)
503520

504-
message_grouper = TestReader(MAX_PAGES_PER_SLICE, MAX_SLICES)
521+
message_grouper = TestReader(
522+
max_pages_per_slice=MAX_PAGES_PER_SLICE,
523+
max_slices=MAX_SLICES,
524+
)
505525

506526
actual_response: StreamRead = message_grouper.run_test_read(
507527
source=mock_source,
@@ -614,7 +634,10 @@ def test_get_grouped_messages_with_many_slices(mock_entrypoint_read: Mock) -> No
614634
),
615635
)
616636

617-
connector_builder_handler = TestReader(MAX_PAGES_PER_SLICE, MAX_SLICES)
637+
connector_builder_handler = TestReader(
638+
max_pages_per_slice=MAX_PAGES_PER_SLICE,
639+
max_slices=MAX_SLICES,
640+
)
618641

619642
stream_read: StreamRead = connector_builder_handler.run_test_read(
620643
source=mock_source,
@@ -658,7 +681,10 @@ def test_get_grouped_messages_given_maximum_number_of_slices_then_test_read_limi
658681
),
659682
)
660683

661-
api = TestReader(MAX_PAGES_PER_SLICE, MAX_SLICES)
684+
api = TestReader(
685+
max_pages_per_slice=MAX_PAGES_PER_SLICE,
686+
max_slices=MAX_SLICES,
687+
)
662688

663689
stream_read: StreamRead = api.run_test_read(
664690
source=mock_source,
@@ -686,7 +712,10 @@ def test_get_grouped_messages_given_maximum_number_of_pages_then_test_read_limit
686712
),
687713
)
688714

689-
api = TestReader(MAX_PAGES_PER_SLICE, MAX_SLICES)
715+
api = TestReader(
716+
max_pages_per_slice=MAX_PAGES_PER_SLICE,
717+
max_slices=MAX_SLICES,
718+
)
690719

691720
stream_read: StreamRead = api.run_test_read(
692721
source=mock_source,
@@ -705,7 +734,10 @@ def test_read_stream_returns_error_if_stream_does_not_exist() -> None:
705734

706735
full_config: Mapping[str, Any] = {**CONFIG, **{"__injected_declarative_manifest": MANIFEST}}
707736

708-
message_grouper = TestReader(MAX_PAGES_PER_SLICE, MAX_SLICES)
737+
message_grouper = TestReader(
738+
max_pages_per_slice=MAX_PAGES_PER_SLICE,
739+
max_slices=MAX_SLICES,
740+
)
709741
actual_response = message_grouper.run_test_read(
710742
source=mock_source,
711743
config=full_config,
@@ -730,7 +762,10 @@ def test_given_control_message_then_stream_read_has_config_update(
730762
+ [connector_configuration_control_message(1, updated_config)]
731763
),
732764
)
733-
connector_builder_handler = TestReader(MAX_PAGES_PER_SLICE, MAX_SLICES)
765+
connector_builder_handler = TestReader(
766+
max_pages_per_slice=MAX_PAGES_PER_SLICE,
767+
max_slices=MAX_SLICES,
768+
)
734769
stream_read: StreamRead = connector_builder_handler.run_test_read(
735770
source=mock_source,
736771
config=CONFIG,
@@ -760,7 +795,10 @@ def test_given_multiple_control_messages_then_stream_read_has_latest_based_on_em
760795
]
761796
),
762797
)
763-
connector_builder_handler = TestReader(MAX_PAGES_PER_SLICE, MAX_SLICES)
798+
connector_builder_handler = TestReader(
799+
max_pages_per_slice=MAX_PAGES_PER_SLICE,
800+
max_slices=MAX_SLICES,
801+
)
764802
stream_read: StreamRead = connector_builder_handler.run_test_read(
765803
source=mock_source,
766804
config=CONFIG,
@@ -788,7 +826,10 @@ def test_given_multiple_control_messages_with_same_timestamp_then_stream_read_ha
788826
]
789827
),
790828
)
791-
connector_builder_handler = TestReader(MAX_PAGES_PER_SLICE, MAX_SLICES)
829+
connector_builder_handler = TestReader(
830+
max_pages_per_slice=MAX_PAGES_PER_SLICE,
831+
max_slices=MAX_SLICES,
832+
)
792833
stream_read: StreamRead = connector_builder_handler.run_test_read(
793834
source=mock_source,
794835
config=CONFIG,
@@ -805,7 +846,10 @@ def test_given_auxiliary_requests_then_return_auxiliary_request(mock_entrypoint_
805846
mock_entrypoint_read,
806847
iter(any_request_and_response_with_a_record() + [auxiliary_request_log_message()]),
807848
)
808-
connector_builder_handler = TestReader(MAX_PAGES_PER_SLICE, MAX_SLICES)
849+
connector_builder_handler = TestReader(
850+
max_pages_per_slice=MAX_PAGES_PER_SLICE,
851+
max_slices=MAX_SLICES,
852+
)
809853
stream_read: StreamRead = connector_builder_handler.run_test_read(
810854
source=mock_source,
811855
config=CONFIG,
@@ -819,7 +863,10 @@ def test_given_auxiliary_requests_then_return_auxiliary_request(mock_entrypoint_
819863
@patch("airbyte_cdk.connector_builder.test_reader.reader.AirbyteEntrypoint.read")
820864
def test_given_no_slices_then_return_empty_slices(mock_entrypoint_read: Mock) -> None:
821865
mock_source = make_mock_source(mock_entrypoint_read, iter([auxiliary_request_log_message()]))
822-
connector_builder_handler = TestReader(MAX_PAGES_PER_SLICE, MAX_SLICES)
866+
connector_builder_handler = TestReader(
867+
max_pages_per_slice=MAX_PAGES_PER_SLICE,
868+
max_slices=MAX_SLICES,
869+
)
823870
stream_read: StreamRead = connector_builder_handler.run_test_read(
824871
source=mock_source,
825872
config=CONFIG,
@@ -845,7 +892,10 @@ def test_given_pk_then_ensure_pk_is_pass_to_schema_inferrence(mock_entrypoint_re
845892
mock_source.streams.return_value = [Mock()]
846893
mock_source.streams.return_value[0].primary_key = [["id"]]
847894
mock_source.streams.return_value[0].cursor_field = _NO_CURSOR_FIELD
848-
connector_builder_handler = TestReader(MAX_PAGES_PER_SLICE, MAX_SLICES)
895+
connector_builder_handler = TestReader(
896+
max_pages_per_slice=MAX_PAGES_PER_SLICE,
897+
max_slices=MAX_SLICES,
898+
)
849899

850900
stream_read: StreamRead = connector_builder_handler.run_test_read(
851901
source=mock_source,
@@ -874,7 +924,10 @@ def test_given_cursor_field_then_ensure_cursor_field_is_pass_to_schema_inferrenc
874924
mock_source.streams.return_value = [Mock()]
875925
mock_source.streams.return_value[0].primary_key = _NO_PK
876926
mock_source.streams.return_value[0].cursor_field = ["date"]
877-
connector_builder_handler = TestReader(MAX_PAGES_PER_SLICE, MAX_SLICES)
927+
connector_builder_handler = TestReader(
928+
max_pages_per_slice=MAX_PAGES_PER_SLICE,
929+
max_slices=MAX_SLICES,
930+
)
878931

879932
stream_read: StreamRead = connector_builder_handler.run_test_read(
880933
source=mock_source,

unit_tests/sources/declarative/test_yaml_declarative_source.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -86,7 +86,7 @@ def test_source_is_created_if_toplevel_fields_are_known(self):
8686
type: CheckStream
8787
stream_names: ["lists"]
8888
"""
89-
temporary_file = TestFileContent(content)
89+
temporary_file = TestFileContent(content=content)
9090
MockYamlDeclarativeSource(temporary_file.filename)
9191

9292
def test_source_fails_for_invalid_yaml(self):
@@ -104,7 +104,7 @@ def test_source_fails_for_invalid_yaml(self):
104104
type: CheckStream
105105
stream_names: ["lists"]
106106
"""
107-
temporary_file = TestFileContent(content)
107+
temporary_file = TestFileContent(content=content)
108108
with pytest.raises(ParserError):
109109
MockYamlDeclarativeSource(temporary_file.filename)
110110

@@ -127,7 +127,7 @@ def test_source_with_missing_reference_fails(self):
127127
type: CheckStream
128128
stream_names: ["lists"]
129129
"""
130-
temporary_file = TestFileContent(content)
130+
temporary_file = TestFileContent(content=content)
131131
with pytest.raises(UndefinedReferenceException):
132132
MockYamlDeclarativeSource(temporary_file.filename)
133133

unit_tests/sources/file_based/scenarios/check_scenarios.py

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -140,7 +140,8 @@
140140
.set_source_builder(
141141
_base_failure_scenario.source_builder.copy().set_stream_reader(
142142
TestErrorListMatchingFilesInMemoryFilesStreamReader(
143-
files=_base_failure_scenario.source_builder._files, file_type="csv"
143+
files=_base_failure_scenario.source_builder._files,
144+
file_type="csv",
144145
)
145146
)
146147
)
@@ -154,7 +155,8 @@
154155
.set_source_builder(
155156
_base_failure_scenario.source_builder.copy().set_stream_reader(
156157
TestErrorOpenFileInMemoryFilesStreamReader(
157-
files=_base_failure_scenario.source_builder._files, file_type="csv"
158+
files=_base_failure_scenario.source_builder._files,
159+
file_type="csv",
158160
)
159161
)
160162
)

unit_tests/sources/file_based/scenarios/scenario_builder.py

Lines changed: 15 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -240,21 +240,21 @@ def build(self) -> "TestScenario[SourceType]":
240240
state,
241241
)
242242
return TestScenario(
243-
self._name,
244-
self._config,
245-
source,
246-
self._expected_spec,
247-
self._expected_check_status,
248-
self._expected_catalog,
249-
self._expected_logs,
250-
self._expected_records,
251-
self._expected_check_error,
252-
self._expected_discover_error,
253-
self._expected_read_error,
254-
self._incremental_scenario_config,
255-
self._expected_analytics,
256-
self._log_levels,
257-
self._catalog,
243+
name=self._name,
244+
config=self._config,
245+
source=source,
246+
expected_spec=self._expected_spec,
247+
expected_check_status=self._expected_check_status,
248+
expected_catalog=self._expected_catalog,
249+
expected_logs=self._expected_logs,
250+
expected_records=self._expected_records,
251+
expected_check_error=self._expected_check_error,
252+
expected_discover_error=self._expected_discover_error,
253+
expected_read_error=self._expected_read_error,
254+
incremental_scenario_config=self._incremental_scenario_config,
255+
expected_analytics=self._expected_analytics,
256+
log_levels=self._log_levels,
257+
catalog=self._catalog,
258258
)
259259

260260
def _configured_catalog(self, sync_mode: SyncMode) -> Optional[Mapping[str, Any]]:

0 commit comments

Comments
 (0)