airbyte-python-cdk/unit_tests/sources/declarative/requesters/query_properties/test_property_chunking.py at f76ee0871128af57e3a9e89f58dbb1e3f682ed5a · airbytehq/airbyte-python-cdk · GitHub

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
# Copyright (c) 2025 Airbyte, Inc., all rights reserved.

import pytest

from airbyte_cdk.sources.declarative.requesters.query_properties import PropertyChunking
from airbyte_cdk.sources.declarative.requesters.query_properties.property_chunking import (
    PropertyLimitType,
)
from airbyte_cdk.sources.declarative.requesters.query_properties.strategies import GroupByKey
from airbyte_cdk.sources.types import Record

CONFIG = {}


@pytest.mark.parametrize(
    "property_fields,always_include_properties,property_limit_type,property_limit,expected_property_chunks",
    [
        pytest.param(
            ["rick", "chelsea", "victoria", "tim", "saxon", "lochlan", "piper"],
            None,
            PropertyLimitType.property_count,
            2,
            [["rick", "chelsea"], ["victoria", "tim"], ["saxon", "lochlan"], ["piper"]],
            id="test_property_chunking",
        ),
        pytest.param(
            ["rick", "chelsea", "victoria", "tim"],
            ["mook", "gaitok"],
            PropertyLimitType.property_count,
            2,
            [["mook", "gaitok", "rick", "chelsea"], ["mook", "gaitok", "victoria", "tim"]],
            id="test_property_chunking_with_always_include_fields",
        ),
        pytest.param(
            ["rick", "chelsea", "victoria", "tim", "saxon", "lochlan", "piper"],
            None,
            PropertyLimitType.property_count,
            None,
            [["rick", "chelsea", "victoria", "tim", "saxon", "lochlan", "piper"]],
            id="test_property_chunking_no_limit",
        ),
        pytest.param(
            ["kate", "laurie", "jaclyn"],
            None,
            PropertyLimitType.characters,
            20,
            [["kate", "laurie"], ["jaclyn"]],
            id="test_property_chunking_limit_characters",
        ),
        pytest.param(
            ["laurie", "jaclyn", "kaitlin"],
            None,
            PropertyLimitType.characters,
            17,  # laurie%2Cjaclyn%2C == 18, so this will create separate chunks
            [["laurie"], ["jaclyn"], ["kaitlin"]],
            id="test_property_chunking_includes_extra_delimiter",
        ),
        pytest.param(
            [],
            None,
            PropertyLimitType.property_count,
            5,
            [[]],
            id="test_property_chunking_no_properties",
        ),
    ],
)
def test_get_request_property_chunks(
    property_fields,
    always_include_properties,
    property_limit_type,
    property_limit,
    expected_property_chunks,
):
    property_chunking = PropertyChunking(
        property_limit_type=property_limit_type,
        property_limit=property_limit,
        record_merge_strategy=GroupByKey(key="id", config=CONFIG, parameters={}),
        config=CONFIG,
        parameters={},
    )

    property_chunks = list(
        property_chunking.get_request_property_chunks(
            property_fields=property_fields, always_include_properties=always_include_properties
        )
    )

    assert len(property_chunks) == len(expected_property_chunks)
    for i, expected_property_chunk in enumerate(expected_property_chunks):
        assert property_chunks[i] == expected_property_chunk


def test_get_merge_key():
    record = Record(stream_name="test", data={"id": "0"})
    property_chunking = PropertyChunking(
        property_limit_type=PropertyLimitType.property_count,
        property_limit=10,
        record_merge_strategy=GroupByKey(key="id", config=CONFIG, parameters={}),
        config=CONFIG,
        parameters={},
    )

    merge_key = property_chunking.get_merge_key(record=record)
    assert merge_key == "0"


def test_given_single_property_chunk_when_get_request_property_chunks_then_always_include_properties_are_not_added_to_input_list():
    """
    This test is used to validate that we don't manipulate the in-memory values from get_request_property_chunks
    """
    property_chunking = PropertyChunking(
        property_limit_type=PropertyLimitType.property_count,
        property_limit=None,
        record_merge_strategy=GroupByKey(key="id", config=CONFIG, parameters={}),
        config=CONFIG,
        parameters={},
    )

    property_fields = ["rick", "chelsea", "victoria", "tim", "saxon", "lochlan", "piper"]
    list(
        property_chunking.get_request_property_chunks(
            property_fields=property_fields,
            always_include_properties=["id"],
        )
    )

    assert property_fields == ["rick", "chelsea", "victoria", "tim", "saxon", "lochlan", "piper"]