Merge remote-tracking branch 'origin/main' into baz/cdk/extract-common-manifest-parts

Author: Oleksandr Bazarnov

airbyte_cdk/sources/declarative/auth/oauth.py

@@ -239,8 +239,8 @@ def get_token_expiry_date(self) -> AirbyteDateTime:
     def _has_access_token_been_initialized(self) -> bool:
         return self._access_token is not None
 
-    def set_token_expiry_date(self, value: Union[str, int]) -> None:
-        self._token_expiry_date = self._parse_token_expiration_date(value)
+    def set_token_expiry_date(self, value: AirbyteDateTime) -> None:
+        self._token_expiry_date = value
 
     def get_assertion_name(self) -> str:
         return self.assertion_name

airbyte_cdk/sources/streams/http/requests_native_auth/abstract_oauth.py

@@ -130,7 +130,7 @@ def build_refresh_request_headers(self) -> Mapping[str, Any] | None:
         headers = self.get_refresh_request_headers()
         return headers if headers else None
 
-    def refresh_access_token(self) -> Tuple[str, Union[str, int]]:
+    def refresh_access_token(self) -> Tuple[str, AirbyteDateTime]:
         """
         Returns the refresh token and its expiration datetime
 
@@ -148,6 +148,14 @@ def refresh_access_token(self) -> Tuple[str, Union[str, int]]:
     # PRIVATE METHODS
     # ----------------
 
+    def _default_token_expiry_date(self) -> AirbyteDateTime:
+        """
+        Returns the default token expiry date
+        """
+        # 1 hour was chosen as a middle ground to avoid unnecessary frequent refreshes and token expiration
+        default_token_expiry_duration_hours = 1  # 1 hour
+        return ab_datetime_now() + timedelta(hours=default_token_expiry_duration_hours)
+
     def _wrap_refresh_token_exception(
         self, exception: requests.exceptions.RequestException
     ) -> bool:
@@ -257,14 +265,10 @@ def _ensure_access_token_in_response(self, response_data: Mapping[str, Any]) ->
 
     def _parse_token_expiration_date(self, value: Union[str, int]) -> AirbyteDateTime:
         """
-        Return the expiration datetime of the refresh token
+        Parse a string or integer token expiration date into a datetime object
 
         :return: expiration datetime
         """
-        if not value and not self.token_has_expired():
-            # No expiry token was provided but the previous one is not expired so it's fine
-            return self.get_token_expiry_date()
-
         if self.token_expiry_is_time_of_expiration:
             if not self.token_expiry_date_format:
                 raise ValueError(
@@ -308,17 +312,30 @@ def _extract_refresh_token(self, response_data: Mapping[str, Any]) -> Any:
         """
         return self._find_and_get_value_from_response(response_data, self.get_refresh_token_name())
 
-    def _extract_token_expiry_date(self, response_data: Mapping[str, Any]) -> Any:
+    def _extract_token_expiry_date(self, response_data: Mapping[str, Any]) -> AirbyteDateTime:
         """
         Extracts the token_expiry_date, like `expires_in` or `expires_at`, etc from the given response data.
 
+        If the token_expiry_date is not found, it will return an existing token expiry date if set, or a default token expiry date.
+
         Args:
             response_data (Mapping[str, Any]): The response data from which to extract the token_expiry_date.
 
         Returns:
-            str: The extracted token_expiry_date.
+            The extracted token_expiry_date or None if not found.
         """
-        return self._find_and_get_value_from_response(response_data, self.get_expires_in_name())
+        expires_in = self._find_and_get_value_from_response(
+            response_data, self.get_expires_in_name()
+        )
+        if expires_in is not None:
+            return self._parse_token_expiration_date(expires_in)
+
+        # expires_in is None
+        existing_expiry_date = self.get_token_expiry_date()
+        if existing_expiry_date and not self.token_has_expired():
+            return existing_expiry_date
+
+        return self._default_token_expiry_date()
 
     def _find_and_get_value_from_response(
         self,
@@ -344,7 +361,7 @@ def _find_and_get_value_from_response(
         """
         if current_depth > max_depth:
             # this is needed to avoid an inf loop, possible with a very deep nesting observed.
-            message = f"The maximum level of recursion is reached. Couldn't find the speficied `{key_name}` in the response."
+            message = f"The maximum level of recursion is reached. Couldn't find the specified `{key_name}` in the response."
             raise ResponseKeysMaxRecurtionReached(
                 internal_message=message, message=message, failure_type=FailureType.config_error
             )
@@ -441,7 +458,7 @@ def get_token_expiry_date(self) -> AirbyteDateTime:
         """Expiration date of the access token"""
 
     @abstractmethod
-    def set_token_expiry_date(self, value: Union[str, int]) -> None:
+    def set_token_expiry_date(self, value: AirbyteDateTime) -> None:
         """Setter for access token expiration date"""
 
     @abstractmethod

airbyte_cdk/sources/streams/http/requests_native_auth/oauth.py

@@ -120,8 +120,8 @@ def get_grant_type(self) -> str:
     def get_token_expiry_date(self) -> AirbyteDateTime:
         return self._token_expiry_date
 
-    def set_token_expiry_date(self, value: Union[str, int]) -> None:
-        self._token_expiry_date = self._parse_token_expiration_date(value)
+    def set_token_expiry_date(self, value: AirbyteDateTime) -> None:
+        self._token_expiry_date = value
 
     @property
     def token_expiry_is_time_of_expiration(self) -> bool:
@@ -316,26 +316,6 @@ def token_has_expired(self) -> bool:
         """Returns True if the token is expired"""
         return ab_datetime_now() > self.get_token_expiry_date()
 
-    @staticmethod
-    def get_new_token_expiry_date(
-        access_token_expires_in: str,
-        token_expiry_date_format: str | None = None,
-    ) -> AirbyteDateTime:
-        """
-        Calculate the new token expiry date based on the provided expiration duration or format.
-
-        Args:
-            access_token_expires_in (str): The duration (in seconds) until the access token expires, or the expiry date in a specific format.
-            token_expiry_date_format (str | None, optional): The format of the expiry date if provided. Defaults to None.
-
-        Returns:
-            AirbyteDateTime: The calculated expiry date of the access token.
-        """
-        if token_expiry_date_format:
-            return ab_datetime_parse(access_token_expires_in)
-        else:
-            return ab_datetime_now() + timedelta(seconds=int(access_token_expires_in))
-
     def get_access_token(self) -> str:
         """Retrieve new access and refresh token if the access token has expired.
         The new refresh token is persisted with the set_refresh_token function
@@ -346,16 +326,13 @@ def get_access_token(self) -> str:
             new_access_token, access_token_expires_in, new_refresh_token = (
                 self.refresh_access_token()
             )
-            new_token_expiry_date: AirbyteDateTime = self.get_new_token_expiry_date(
-                access_token_expires_in, self._token_expiry_date_format
-            )
             self.access_token = new_access_token
             self.set_refresh_token(new_refresh_token)
-            self.set_token_expiry_date(new_token_expiry_date)
+            self.set_token_expiry_date(access_token_expires_in)
             self._emit_control_message()
         return self.access_token
 
-    def refresh_access_token(self) -> Tuple[str, str, str]:  # type: ignore[override]
+    def refresh_access_token(self) -> Tuple[str, AirbyteDateTime, str]:  # type: ignore[override]
         """
         Refreshes the access token by making a handled request and extracting the necessary token information.
 

airbyte_cdk/test/entrypoint_wrapper.py

@@ -82,6 +82,10 @@ def records(self) -> List[AirbyteMessage]:
     def state_messages(self) -> List[AirbyteMessage]:
         return self._get_message_by_types([Type.STATE])
 
+    @property
+    def connection_status_messages(self) -> List[AirbyteMessage]:
+        return self._get_message_by_types([Type.CONNECTION_STATUS])
+
     @property
     def most_recent_state(self) -> Any:
         state_messages = self._get_message_by_types([Type.STATE])

airbyte_cdk/test/standard_tests/__init__.py

@@ -0,0 +1,46 @@
+# Copyright (c) 2024 Airbyte, Inc., all rights reserved.
+'''FAST Airbyte Standard Tests
+
+This module provides a set of base classes for declarative connector test suites.
+The goal of this module is to provide a robust and extensible framework for testing Airbyte
+connectors.
+
+Example usage:
+
+```python
+# `test_airbyte_standards.py`
+from airbyte_cdk.test import standard_tests
+
+pytest_plugins = [
+    "airbyte_cdk.test.standard_tests.pytest_hooks",
+]
+
+
+class TestSuiteSourcePokeAPI(standard_tests.DeclarativeSourceTestSuite):
+    """Test suite for the source."""
+```
+
+Available test suites base classes:
+- `DeclarativeSourceTestSuite`: A test suite for declarative sources.
+- `SourceTestSuiteBase`: A test suite for sources.
+- `DestinationTestSuiteBase`: A test suite for destinations.
+
+'''
+
+from airbyte_cdk.test.standard_tests.connector_base import (
+    ConnectorTestScenario,
+    ConnectorTestSuiteBase,
+)
+from airbyte_cdk.test.standard_tests.declarative_sources import (
+    DeclarativeSourceTestSuite,
+)
+from airbyte_cdk.test.standard_tests.destination_base import DestinationTestSuiteBase
+from airbyte_cdk.test.standard_tests.source_base import SourceTestSuiteBase
+
+__all__ = [
+    "ConnectorTestScenario",
+    "ConnectorTestSuiteBase",
+    "DeclarativeSourceTestSuite",
+    "DestinationTestSuiteBase",
+    "SourceTestSuiteBase",
+]

airbyte_cdk/test/standard_tests/_job_runner.py

@@ -0,0 +1,159 @@
+# Copyright (c) 2025 Airbyte, Inc., all rights reserved.
+"""Job runner for Airbyte Standard Tests."""
+
+import logging
+import tempfile
+import uuid
+from dataclasses import asdict
+from pathlib import Path
+from typing import Any, Callable, Literal
+
+import orjson
+from typing_extensions import Protocol, runtime_checkable
+
+from airbyte_cdk.models import (
+    ConfiguredAirbyteCatalog,
+    Status,
+)
+from airbyte_cdk.test import entrypoint_wrapper
+from airbyte_cdk.test.standard_tests.models import (
+    ConnectorTestScenario,
+)
+
+
+def _errors_to_str(
+    entrypoint_output: entrypoint_wrapper.EntrypointOutput,
+) -> str:
+    """Convert errors from entrypoint output to a string."""
+    if not entrypoint_output.errors:
+        # If there are no errors, return an empty string.
+        return ""
+
+    return "\n" + "\n".join(
+        [
+            str(error.trace.error).replace(
+                "\\n",
+                "\n",
+            )
+            for error in entrypoint_output.errors
+            if error.trace
+        ],
+    )
+
+
+@runtime_checkable
+class IConnector(Protocol):
+    """A connector that can be run in a test scenario.
+
+    Note: We currently use 'spec' to determine if we have a connector object.
+    In the future, it would be preferred to leverage a 'launch' method instead,
+    directly on the connector (which doesn't yet exist).
+    """
+
+    def spec(self, logger: logging.Logger) -> Any:
+        """Connectors should have a `spec` method."""
+
+
+def run_test_job(
+    connector: IConnector | type[IConnector] | Callable[[], IConnector],
+    verb: Literal["read", "check", "discover"],
+    test_scenario: ConnectorTestScenario,
+    *,
+    catalog: ConfiguredAirbyteCatalog | dict[str, Any] | None = None,
+) -> entrypoint_wrapper.EntrypointOutput:
+    """Run a test scenario from provided CLI args and return the result."""
+    if not connector:
+        raise ValueError("Connector is required")
+
+    if catalog and isinstance(catalog, ConfiguredAirbyteCatalog):
+        # Convert the catalog to a dict if it's already a ConfiguredAirbyteCatalog.
+        catalog = asdict(catalog)
+
+    connector_obj: IConnector
+    if isinstance(connector, type) or callable(connector):
+        # If the connector is a class or a factory lambda, instantiate it.
+        connector_obj = connector()
+    elif isinstance(connector, IConnector):
+        connector_obj = connector
+    else:
+        raise ValueError(
+            f"Invalid connector input: {type(connector)}",
+        )
+
+    args: list[str] = [verb]
+    if test_scenario.config_path:
+        args += ["--config", str(test_scenario.config_path)]
+    elif test_scenario.config_dict:
+        config_path = (
+            Path(tempfile.gettempdir()) / "airbyte-test" / f"temp_config_{uuid.uuid4().hex}.json"
+        )
+        config_path.parent.mkdir(parents=True, exist_ok=True)
+        config_path.write_text(orjson.dumps(test_scenario.config_dict).decode())
+        args += ["--config", str(config_path)]
+
+    catalog_path: Path | None = None
+    if verb not in ["discover", "check"]:
+        # We need a catalog for read.
+        if catalog:
+            # Write the catalog to a temp json file and pass the path to the file as an argument.
+            catalog_path = (
+                Path(tempfile.gettempdir())
+                / "airbyte-test"
+                / f"temp_catalog_{uuid.uuid4().hex}.json"
+            )
+            catalog_path.parent.mkdir(parents=True, exist_ok=True)
+            catalog_path.write_text(orjson.dumps(catalog).decode())
+        elif test_scenario.configured_catalog_path:
+            catalog_path = Path(test_scenario.configured_catalog_path)
+
+        if catalog_path:
+            args += ["--catalog", str(catalog_path)]
+
+    # This is a bit of a hack because the source needs the catalog early.
+    # Because it *also* can fail, we have to redundantly wrap it in a try/except block.
+
+    result: entrypoint_wrapper.EntrypointOutput = entrypoint_wrapper._run_command(  # noqa: SLF001  # Non-public API
+        source=connector_obj,  # type: ignore [arg-type]
+        args=args,
+        expecting_exception=test_scenario.expect_exception,
+    )
+    if result.errors and not test_scenario.expect_exception:
+        raise AssertionError(
+            f"Expected no errors but got {len(result.errors)}: \n" + _errors_to_str(result)
+        )
+
+    if verb == "check":
+        # Check is expected to fail gracefully without an exception.
+        # Instead, we assert that we have a CONNECTION_STATUS message with
+        # a failure status.
+        assert len(result.connection_status_messages) == 1, (
+            "Expected exactly one CONNECTION_STATUS message. Got "
+            f"{len(result.connection_status_messages)}:\n"
+            + "\n".join([str(msg) for msg in result.connection_status_messages])
+            + _errors_to_str(result)
+        )
+        if test_scenario.expect_exception:
+            conn_status = result.connection_status_messages[0].connectionStatus
+            assert conn_status, (
+                "Expected CONNECTION_STATUS message to be present. Got: \n"
+                + "\n".join([str(msg) for msg in result.connection_status_messages])
+            )
+            assert conn_status.status == Status.FAILED, (
+                "Expected CONNECTION_STATUS message to be FAILED. Got: \n"
+                + "\n".join([str(msg) for msg in result.connection_status_messages])
+            )
+
+        return result
+
+    # For all other verbs, we assert check that an exception is raised (or not).
+    if test_scenario.expect_exception:
+        if not result.errors:
+            raise AssertionError("Expected exception but got none.")
+
+        return result
+
+    assert not result.errors, (
+        f"Expected no errors but got {len(result.errors)}: \n" + _errors_to_str(result)
+    )
+
+    return result