Adding tests and fixing some stuff

Author: maxi297

airbyte_cdk/sources/declarative/schema/dynamic_schema_loader.py

@@ -114,7 +114,28 @@ def _update_pointer(
 
 
 @deprecated("This class is experimental. Use at your own risk.", category=ExperimentalClassWarning)
-class AdditionalPropertyFieldsInferer(ABC):
+class AdditionalPropertyFieldsInferrer(ABC):
+    """
+    Infers additional fields to be added to each property. For example, if this inferrer returns {"toto": "tata"}, a property that would have looked like this:
+    ```
+        "properties": {
+            "Id": {
+                "type": ["null", "string"],
+            },
+            <...>
+        }
+    ```
+    ... will look like this:
+        ```
+        "properties": {
+            "Id": {
+                "type": ["null", "string"],
+                "toto": "tata"
+            },
+            <...>
+        }
+    ```
+    """
     @abstractmethod
     def infer(self, property_definition: MutableMapping[str, Any]) -> MutableMapping[str, Any]:
         """
@@ -135,7 +156,8 @@ class DynamicSchemaLoader(SchemaLoader):
     parameters: InitVar[Mapping[str, Any]]
     schema_type_identifier: SchemaTypeIdentifier
     schema_transformations: List[RecordTransformation] = field(default_factory=lambda: [])
-    additional_property_fields_inferrer: Optional[AdditionalPropertyFieldsInferer] = None
+    additional_property_fields_inferrer: Optional[AdditionalPropertyFieldsInferrer] = None
+    allow_additional_properties: bool = True
 
     def get_json_schema(self) -> Mapping[str, Any]:
         """
@@ -168,7 +190,7 @@ def get_json_schema(self) -> Mapping[str, Any]:
         return {
             "$schema": "https://json-schema.org/draft-07/schema#",
             "type": "object",
-            "additionalProperties": True,
+            "additionalProperties": self.allow_additional_properties,
             "properties": transformed_properties,
         }
 

airbyte_cdk/test/catalog_builder.py

@@ -2,6 +2,8 @@
 
 from typing import Any, Dict, List, Union, overload
 
+from airbyte_protocol_dataclasses.models import DestinationSyncMode
+
 from airbyte_cdk.models import (
     ConfiguredAirbyteCatalog,
     ConfiguredAirbyteStream,
@@ -32,6 +34,10 @@ def with_sync_mode(self, sync_mode: SyncMode) -> "ConfiguredAirbyteStreamBuilder
         self._stream["sync_mode"] = sync_mode.name
         return self
 
+    def with_destination_sync_mode(self, sync_mode: DestinationSyncMode) -> "ConfiguredAirbyteStreamBuilder":
+        self._stream["destination_sync_mode"] = sync_mode.name
+        return self
+
     def with_primary_key(self, pk: List[List[str]]) -> "ConfiguredAirbyteStreamBuilder":
         self._stream["primary_key"] = pk
         self._stream["stream"]["source_defined_primary_key"] = pk  # type: ignore  # we assume that self._stream["stream"] is a Dict[str, Any]

airbyte_cdk/test/mock_http/request.py

@@ -72,7 +72,11 @@ def _to_mapping(
         elif isinstance(body, bytes):
             return json.loads(body.decode())  # type: ignore  # assumes return type of Mapping[str, Any]
         elif isinstance(body, str):
-            return json.loads(body)  # type: ignore  # assumes return type of Mapping[str, Any]
+            try:
+                return json.loads(body)
+            except json.JSONDecodeError:
+                # one of the body is a mapping while the other isn't so comparison should fail anyway
+                return None
         return None
 
     @staticmethod

unit_tests/sources/declarative/schema/test_dynamic_schema_loader.py

@@ -4,7 +4,8 @@
 
 import json
 from copy import deepcopy
-from unittest.mock import MagicMock
+from typing import Any, Mapping, MutableMapping
+from unittest.mock import MagicMock, Mock
 
 import pytest
 
@@ -14,13 +15,21 @@
 from airbyte_cdk.sources.declarative.parsers.model_to_component_factory import (
     ModelToComponentFactory,
 )
-from airbyte_cdk.sources.declarative.schema import DynamicSchemaLoader, SchemaTypeIdentifier
+from airbyte_cdk.sources.declarative.retrievers import Retriever
+from airbyte_cdk.sources.declarative.schema import (
+    DynamicSchemaLoader,
+    SchemaTypeIdentifier,
+    TypesMap,
+)
+from airbyte_cdk.sources.declarative.schema.dynamic_schema_loader import (
+    AdditionalPropertyFieldsInferrer,
+)
 from airbyte_cdk.test.mock_http import HttpMocker, HttpRequest, HttpResponse
 
 _CONFIG = {
     "start_date": "2024-07-01T00:00:00.000Z",
 }
-
+_ANY_PARAMETERS = {}
 _MANIFEST = {
     "version": "6.7.0",
     "definitions": {
@@ -412,3 +421,87 @@ def test_dynamic_schema_loader_with_type_conditions():
 
     assert len(actual_catalog.streams) == 1
     assert actual_catalog.streams[0].json_schema == expected_schema
+
+
+def _mock_schema_loader_retriever(http_response_body) -> Retriever:
+    retriever = Mock(spec=Retriever)
+    retriever.read_records.return_value = iter([http_response_body])
+    return retriever
+
+
+class TestAdditionalPropertyFieldsInferrer(AdditionalPropertyFieldsInferrer):
+    def __init__(self, properties_to_add: Mapping[str, Any]):
+        self._properties_to_add = properties_to_add
+
+    def infer(self, property_definition: MutableMapping[str, Any]) -> Mapping[str, Any]:
+        return self._properties_to_add
+
+
+def test_additional_property_fields_inferrer():
+    properties_to_add = {"added_property": "a_value"}
+    expected_schema = {
+        "$schema": "https://json-schema.org/draft-07/schema#",
+        "additionalProperties": False,
+        "type": "object",
+        "properties": {
+            "id": {"type": ["null", "integer"]} | properties_to_add,
+        },
+    }
+    schema_loader = DynamicSchemaLoader(
+        retriever=_mock_schema_loader_retriever({"fields": [{"name": "id", "type": "integer"}]}),
+        additional_property_fields_inferrer=TestAdditionalPropertyFieldsInferrer(properties_to_add),
+        schema_type_identifier=SchemaTypeIdentifier(
+            key_pointer=["name"],
+            type_pointer=["type"],
+            types_mapping=[
+                TypesMap(
+                    current_type="integer",
+                    target_type="integer",
+                    condition=None,
+                ),
+            ],
+            schema_pointer=["fields"],
+            parameters=_ANY_PARAMETERS,
+        ),
+        allow_additional_properties=False,
+        config={},
+        parameters=_ANY_PARAMETERS,
+    )
+
+    schema = schema_loader.get_json_schema()
+
+    assert schema == expected_schema
+
+
+def test_additional_properties():
+    expected_schema = {
+        "$schema": "https://json-schema.org/draft-07/schema#",
+        "additionalProperties": False,
+        "type": "object",
+        "properties": {
+            "id": {"type": ["null", "integer"]},
+        },
+    }
+    schema_loader = DynamicSchemaLoader(
+        retriever=_mock_schema_loader_retriever({"fields": [{"name": "id", "type": "integer"}]}),
+        schema_type_identifier=SchemaTypeIdentifier(
+            key_pointer=["name"],
+            type_pointer=["type"],
+            types_mapping=[
+                TypesMap(
+                    current_type="integer",
+                    target_type="integer",
+                    condition=None,
+                ),
+            ],
+            schema_pointer=["fields"],
+            parameters=_ANY_PARAMETERS,
+        ),
+        allow_additional_properties=False,
+        config={},
+        parameters=_ANY_PARAMETERS,
+    )
+
+    schema = schema_loader.get_json_schema()
+
+    assert schema == expected_schema

unit_tests/test/mock_http/test_request.py

@@ -156,3 +156,8 @@ def test_given_any_matcher_for_both_when_matches_then_return_true(self):
         request_to_match = HttpRequest("mock://test.com/path", ANY_QUERY_PARAMS)
         actual_request = HttpRequest("mock://test.com/path", ANY_QUERY_PARAMS)
         assert actual_request.matches(request_to_match)
+
+    def test_given_on_match_is_mapping_but_not_input_when_matches_then_return_false(self):
+        request_to_match = HttpRequest("mock://test.com/path", body={"first_field": "another value"})
+        actual_request = HttpRequest("mock://test.com/path", body="another_request_body")
+        assert not actual_request.matches(request_to_match)