Fix mypy and unittest

lazebnyi · lazebnyi · commit 706ba21288a3 · 2025-03-28T17:28:35.000+01:00
diff --git a/airbyte_cdk/sources/declarative/checks/check_stream.py b/airbyte_cdk/sources/declarative/checks/check_stream.py
@@ -5,7 +5,7 @@
 import logging
 import traceback
 from dataclasses import InitVar, dataclass
-from typing import Any, List, Mapping, Tuple
+from typing import Any, List, Mapping, Tuple, Dict, Optional
 
 from airbyte_cdk import AbstractSource
 from airbyte_cdk.sources.declarative.checks.connection_checker import ConnectionChecker
@@ -32,81 +32,105 @@ class CheckStream(ConnectionChecker):
     """
 
     stream_names: List[str]
-    dynamic_streams_check_configs: List[DynamicStreamCheckConfig]
     parameters: InitVar[Mapping[str, Any]]
+    dynamic_streams_check_configs: Optional[List[DynamicStreamCheckConfig]] = None
 
     def __post_init__(self, parameters: Mapping[str, Any]) -> None:
         self._parameters = parameters
+        if self.dynamic_streams_check_configs is None:
+            self.dynamic_streams_check_configs = []
+
+    def _log_error(self, logger: logging.Logger, action: str, error: Exception) -> Tuple[bool, str]:
+        """Logs an error and returns a formatted error message."""
+        error_message = f"Encountered an error while {action}. Error: {error}"
+        logger.error(error_message + f"Error traceback: \n {traceback.format_exc()}", exc_info=True)
+        return False, error_message
 
     def check_connection(
-        self, source: AbstractSource, logger: logging.Logger, config: Mapping[str, Any]
+            self, source: AbstractSource, logger: logging.Logger, config: Mapping[str, Any]
     ) -> Tuple[bool, Any]:
+        """Checks the connection to the source and its streams."""
         try:
             streams = source.streams(config=config)
+            if not streams:
+                return False, f"No streams to connect to from source {source}"
         except Exception as error:
-            error_message = (
-                f"Encountered an error trying to connect to streams. Error: {error}"
-            )
-            logger.error(error_message, exc_info=True)
-            return False, error_message
+            return self._log_error(logger, "connecting to streams", error)
 
         stream_name_to_stream = {s.name: s for s in streams}
-        if len(streams) == 0:
-            return False, f"No streams to connect to from source {source}"
         for stream_name in self.stream_names:
-            if stream_name not in stream_name_to_stream.keys():
+            if stream_name not in stream_name_to_stream:
                 raise ValueError(
-                    f"{stream_name} is not part of the catalog. Expected one of {stream_name_to_stream.keys()}."
+                    f"{stream_name} is not part of the catalog. Expected one of {list(stream_name_to_stream.keys())}."
                 )
 
+            stream_availability, message = self._check_stream_availability(stream_name_to_stream, stream_name, logger)
+            if not stream_availability:
+                return stream_availability, message
+
+        should_check_dynamic_streams = hasattr(source, "resolved_manifest") and hasattr(source, "dynamic_streams") and self.dynamic_streams_check_configs
+
+        if should_check_dynamic_streams:
+            return self._check_dynamic_streams_availability(source, stream_name_to_stream, logger)
+
+        return True, None
+
+    def _check_stream_availability(self, stream_name_to_stream: Dict[str, Any], stream_name: str, logger: logging.Logger) -> Tuple[bool, Any]:
+        """Checks if streams are available."""
+        availability_strategy = HttpAvailabilityStrategy()
+        try:
             stream = stream_name_to_stream[stream_name]
-            availability_strategy = HttpAvailabilityStrategy()
+            stream_is_available, reason = availability_strategy.check_availability(stream, logger)
+            if not stream_is_available:
+                message = f"Stream {stream_name} is not available: {reason}"
+                logger.warning(message)
+                return stream_is_available, message
+        except Exception as error:
+            return self._log_error(logger, f"checking availability of stream {stream_name}", error)
+        return True, None
+
+    def _check_dynamic_streams_availability(
+            self, source: AbstractSource, stream_name_to_stream: Dict[str, Any], logger: logging.Logger
+    ) -> Tuple[bool, Any]:
+        """Checks the availability of dynamic streams."""
+        dynamic_streams = source.resolved_manifest.get("dynamic_streams", [])  # type: ignore[attr-defined] # The source's resolved_manifest manifest is checked before calling this method
+        dynamic_stream_name_to_dynamic_stream = {
+            ds.get("name", f"dynamic_stream_{i}"): ds for i, ds in enumerate(dynamic_streams)
+        }
+        generated_streams = self._map_generated_streams(source.dynamic_streams) # type: ignore[attr-defined] # The source's dynamic_streams manifest is checked before calling this method
+
+        for check_config in self.dynamic_streams_check_configs: # type: ignore[union-attr] # None value for self.dynamic_streams_check_configs handled in __post_init__
+            if check_config.dynamic_stream_name not in dynamic_stream_name_to_dynamic_stream:
+                return False, f"Dynamic stream {check_config.dynamic_stream_name} is not found in manifest."
+
+            generated = generated_streams.get(check_config.dynamic_stream_name, [])
+            stream_availability, message = self._check_generated_streams_availability(generated, stream_name_to_stream, logger, check_config.stream_count)
+            if not stream_availability:
+                return stream_availability, message
+
+        return True, None
+
+    def _map_generated_streams(self, dynamic_streams: List[Dict[str, Any]]) -> Dict[str, List[Dict[str, Any]]]:
+        """Maps dynamic stream names to their corresponding generated streams."""
+        mapped_streams: Dict[str, List[Dict[str, Any]]] = {}
+        for stream in dynamic_streams:
+            mapped_streams.setdefault(stream["dynamic_stream_name"], []).append(stream)
+        return mapped_streams
+
+    def _check_generated_streams_availability(
+            self, generated_streams: List[Dict[str, Any]], stream_name_to_stream: Dict[str, Any],
+            logger: logging.Logger, max_count: int
+    ) -> Tuple[bool, Any]:
+        """Checks availability of generated dynamic streams."""
+        availability_strategy = HttpAvailabilityStrategy()
+        for declarative_stream in generated_streams[: min(max_count, len(generated_streams))]:
+            stream = stream_name_to_stream[declarative_stream["name"]]
             try:
-                stream_is_available, reason = availability_strategy.check_availability(
-                    stream, logger
-                )
+                stream_is_available, reason = availability_strategy.check_availability(stream, logger)
                 if not stream_is_available:
-                    return False, reason
+                    message = f"Dynamic Stream {stream.name} is not available: {reason}"
+                    logger.warning(message)
+                    return False, message
             except Exception as error:
-                logger.error(
-                    f"Encountered an error trying to connect to stream {stream_name}. Error: \n {traceback.format_exc()}"
-                )
-                return False, f"Unable to connect to stream {stream_name} - {error}"
-
-        if hasattr(source, "resolved_manifest") and hasattr(source, "dynamic_streams") and self.dynamic_streams_check_configs:
-            dynamic_stream_name_to_dynamic_stream = {dynamic_stream.get("name", f"dynamic_stream_{i}"): dynamic_stream for i, dynamic_stream in enumerate(source.resolved_manifest.get("dynamic_streams", []))}
-
-            dynamic_stream_name_to_generated_streams = {}
-            for stream in source.dynamic_streams:
-                dynamic_stream_name_to_generated_streams[
-                    stream["dynamic_stream_name"]] = dynamic_stream_name_to_generated_streams.setdefault(
-                    stream["dynamic_stream_name"], []) + [stream]
-
-            for dynamic_streams_check_config in self.dynamic_streams_check_configs:
-                dynamic_stream = dynamic_stream_name_to_dynamic_stream.get(dynamic_streams_check_config.dynamic_stream_name)
-
-                is_config_depend = dynamic_stream["components_resolver"]["type"] == "ConfigComponentsResolver"
-
-                if not is_config_depend and not bool(dynamic_streams_check_config.stream_count):
-                    continue
-
-                generated_streams = dynamic_stream_name_to_generated_streams.get(dynamic_streams_check_config.dynamic_stream_name)
-                availability_strategy = HttpAvailabilityStrategy()
-
-                for declarative_stream in generated_streams[: min(dynamic_streams_check_config.stream_count, len(generated_streams))]:
-                    stream = stream_name_to_stream.get(declarative_stream["name"])
-                    try:
-                        stream_is_available, reason = availability_strategy.check_availability(
-                            stream, logger
-                        )
-                        if not stream_is_available:
-                            logger.warning(f"Stream {stream.name} is not available: {reason}")
-                            return False, reason
-                    except Exception as error:
-                        error_message = (
-                            f"Encountered an error trying to connect to stream {stream.name}. Error: {error}"
-                        )
-                        logger.error(error_message, exc_info=True)
-                        return False, error_message
-
+                return self._log_error(logger, f"checking availability of dynamic stream {stream.name}", error)
         return True, None
diff --git a/unit_tests/sources/declarative/checks/test_check_stream.py b/unit_tests/sources/declarative/checks/test_check_stream.py
@@ -11,6 +11,7 @@
 import pytest
 import requests
 
+from jsonschema.exceptions import ValidationError
 from airbyte_cdk.sources.declarative.checks.check_stream import CheckStream
 from airbyte_cdk.sources.streams.http import HttpStream
 from airbyte_cdk.sources.declarative.concurrent_declarative_source import (
@@ -343,52 +344,97 @@ def parse_response(self, response: requests.Response, **kwargs) -> Iterable[Mapp
             "name": "static_stream",
             "primary_key": "id",
             "schema_loader": {
-                    "type": "InlineSchemaLoader",
-                    "schema": {
-                        "$schema": "http://json-schema.org/schema#",
-                        "properties": {
-                            "id": {"type": "integer"},
-                            "name": {"type": "string"},
-                        },
-                        "type": "object",
+                "type": "InlineSchemaLoader",
+                "schema": {
+                    "$schema": "http://json-schema.org/schema#",
+                    "properties": {
+                        "id": {"type": "integer"},
+                        "name": {"type": "string"},
                     },
-                }
+                    "type": "object",
+                },
+            }
         }
     ]
 }
 
 
 @pytest.mark.parametrize(
-    "check_component",
+    "check_component, expected_result, expectation, response_code, expected_messages",
     [
-        pytest.param({"check": {"type": "CheckStream", "stream_names": ["static_stream"]}},
+        pytest.param({"check": {"type": "CheckStream", "stream_names": ["static_stream"]}}, True, False, 200,
+                     [{"id": 1, "name": "static_1"}, {"id": 2, "name": "static_2"}],
                      id="test_check_only_static_streams"),
         pytest.param({"check": {"type": "CheckStream", "stream_names": ["static_stream"],
                                 "dynamic_streams_check_configs": [
                                     {"type": "DynamicStreamCheckConfig", "dynamic_stream_name": "http_dynamic_stream",
-                                     "stream_count": 1}]}}, id="test_check_static_streams_and_http_dynamic_stream"),
+                                     "stream_count": 1}]}}, True, False, 200, [],
+                     id="test_check_static_streams_and_http_dynamic_stream"),
         pytest.param({"check": {"type": "CheckStream", "stream_names": ["static_stream"],
                                 "dynamic_streams_check_configs": [
                                     {"type": "DynamicStreamCheckConfig", "dynamic_stream_name": "dynamic_stream_1",
-                                     "stream_count": 1}]}}, id="test_check_static_streams_and_config_dynamic_stream"),
-        pytest.param({"check": {"type": "CheckStream", "dynamic_streams_check_configs": [{"type": "DynamicStreamCheckConfig", "dynamic_stream_name": "dynamic_stream_1",
-                                     "stream_count": 1}, {"type": "DynamicStreamCheckConfig", "dynamic_stream_name": "http_dynamic_stream"}]}},
+                                     "stream_count": 1}]}}, True, False, 200, [],
+                     id="test_check_static_streams_and_config_dynamic_stream"),
+        pytest.param({"check": {"type": "CheckStream", "dynamic_streams_check_configs": [
+            {"type": "DynamicStreamCheckConfig", "dynamic_stream_name": "dynamic_stream_1",
+             "stream_count": 1}, {"type": "DynamicStreamCheckConfig", "dynamic_stream_name": "http_dynamic_stream"}]}},
+                     True, False, 200, [],
                      id="test_check_http_dynamic_stream_and_config_dynamic_stream"),
         pytest.param({"check": {"type": "CheckStream", "stream_names": ["static_stream"],
                                 "dynamic_streams_check_configs": [
                                     {"type": "DynamicStreamCheckConfig", "dynamic_stream_name": "dynamic_stream_1",
                                      "stream_count": 1},
-                                    {"type": "DynamicStreamCheckConfig", "dynamic_stream_name": "http_dynamic_stream"}]}},
+                                    {"type": "DynamicStreamCheckConfig",
+                                     "dynamic_stream_name": "http_dynamic_stream"}]}}, True, False, 200, [],
                      id="test_check_static_streams_and_http_dynamic_stream_and_config_dynamic_stream"),
+        pytest.param(
+            {"check": {"type": "CheckStream", "stream_names": ["non_existent_stream"]}},
+            False,
+            True, 200, [],
+            id="test_non_existent_static_stream"
+        ),
+        pytest.param(
+            {"check": {"type": "CheckStream", "dynamic_streams_check_configs": [
+                {"type": "DynamicStreamCheckConfig", "dynamic_stream_name": "unknown_dynamic_stream",
+                 "stream_count": 1}]}}
+            ,
+            False,
+            False, 200, [],
+            id="test_non_existent_dynamic_stream"
+        ),
+        pytest.param({"check": {"type": "CheckStream", "stream_names": ["static_stream"]}}, False, False, 404,
+                     ["Not found. The requested resource was not found on the server."],
+                     id="test_stream_unavailable_unhandled_error"),
+        pytest.param({"check": {"type": "CheckStream", "stream_names": ["static_stream"]}}, False, False, 403,
+                     ["Forbidden. You don't have permission to access this resource."],
+                     id="test_stream_unavailable_handled_error"),
+        pytest.param({"check": {"type": "CheckStream", "stream_names": ["static_stream"]}}, False, False, 401,
+                     ["Unauthorized. Please ensure you are authenticated correctly."],
+                     id="test_stream_unauthorized_error"),
+        pytest.param({"check": {"type": "CheckStream", "dynamic_streams_check_configs": [
+            {"type": "DynamicStreamCheckConfig", "dynamic_stream_name": "dynamic_stream_1",
+             "stream_count": 1}, {"type": "DynamicStreamCheckConfig", "dynamic_stream_name": "http_dynamic_stream"}]}},
+                     False, False, 404, ["Not found. The requested resource was not found on the server."],
+                     id="test_dynamic_stream_unavailable_unhandled_error"),
+        pytest.param({"check": {"type": "CheckStream", "dynamic_streams_check_configs": [
+            {"type": "DynamicStreamCheckConfig", "dynamic_stream_name": "dynamic_stream_1",
+             "stream_count": 1}, {"type": "DynamicStreamCheckConfig", "dynamic_stream_name": "http_dynamic_stream"}]}},
+                     False, False, 403, ["Forbidden. You don't have permission to access this resource."],
+                     id="test_dynamic_stream_unavailable_handled_error"),
+        pytest.param({"check": {"type": "CheckStream", "dynamic_streams_check_configs": [
+            {"type": "DynamicStreamCheckConfig", "dynamic_stream_name": "dynamic_stream_1",
+             "stream_count": 1}, {"type": "DynamicStreamCheckConfig", "dynamic_stream_name": "http_dynamic_stream"}]}},
+                     False, False, 401, ["Unauthorized. Please ensure you are authenticated correctly."],
+                     id="test_dynamic_stream_unauthorized_error"),
     ],
 )
-def test_check_stream(check_component):
+def test_check_stream1(check_component, expected_result, expectation, response_code, expected_messages):
     manifest = {**deepcopy(_MANIFEST_WITHOUT_CHECK_COMPONENT), **check_component}
 
     with HttpMocker() as http_mocker:
         static_stream_request = HttpRequest(url="https://api.test.com/static")
         static_stream_response = HttpResponse(
-            body=json.dumps([{"id": 1, "name": "static_1"}, {"id": 2, "name": "static_2"}])
+            body=json.dumps(expected_messages), status_code=response_code
         )
         http_mocker.get(static_stream_request, static_stream_response)
 
@@ -399,11 +445,11 @@ def test_check_stream(check_component):
         http_mocker.get(items_request, items_response)
 
         item_request = HttpRequest(url="https://api.test.com/items/1")
-        item_response = HttpResponse(body=json.dumps([]), status_code=200)
+        item_response = HttpResponse(body=json.dumps(expected_messages), status_code=response_code)
         http_mocker.get(item_request, item_response)
 
         item_request = HttpRequest(url="https://api.test.com/items/3")
-        item_response = HttpResponse(body=json.dumps([]), status_code=200)
+        item_response = HttpResponse(body=json.dumps(expected_messages), status_code=response_code)
         http_mocker.get(item_request, item_response)
 
         source = ConcurrentDeclarativeSource(
@@ -412,14 +458,32 @@ def test_check_stream(check_component):
             catalog=None,
             state=None,
         )
+        if expectation:
+            with pytest.raises(ValueError):
+                source.check_connection(logger, _CONFIG)
+        else:
+            stream_is_available, reason = source.check_connection(logger, _CONFIG)
 
-        stream_is_available, reason = source.check_connection(logger, _CONFIG)
+            assert stream_is_available == expected_result
 
-    assert stream_is_available
+
+def test_check_stream_missing_fields():
+    """Test if ValueError is raised when dynamic_streams_check_configs is missing required fields."""
+    manifest = {**deepcopy(_MANIFEST_WITHOUT_CHECK_COMPONENT),
+                **{"check": {"type": "CheckStream",
+                             "dynamic_streams_check_configs": [{"type": "DynamicStreamCheckConfig"}]}}}
+    with pytest.raises(ValidationError):
+        source = ConcurrentDeclarativeSource(
+            source_config=manifest,
+            config=_CONFIG,
+            catalog=None,
+            state=None,
+        )
 
 
 def test_check_stream_only_type_provided():
-    manifest = {**deepcopy(_MANIFEST_WITHOUT_CHECK_COMPONENT), **{"check": {"type": "CheckStream"}}}
+    manifest = {**deepcopy(_MANIFEST_WITHOUT_CHECK_COMPONENT),
+                **{"check": {"type": "CheckStream"}}}
     source = ConcurrentDeclarativeSource(
         source_config=manifest,
         config=_CONFIG,
