airbyte/airbyte-cdk/python/unit_tests/sources/fixtures/source_test_fixture.py

#
# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
#

import json
import logging
from abc import ABC
from typing import Any, Iterable, List, Mapping, Optional, Tuple, Union

import requests
from airbyte_cdk.models import (
    AirbyteStream,
    ConfiguredAirbyteCatalog,
    ConfiguredAirbyteStream,
    ConnectorSpecification,
    DestinationSyncMode,
    SyncMode,
)
from airbyte_cdk.sources import AbstractSource
from airbyte_cdk.sources.streams import Stream
from airbyte_cdk.sources.streams.http import HttpStream
from airbyte_cdk.sources.streams.http.requests_native_auth import Oauth2Authenticator
from requests.auth import AuthBase


class SourceTestFixture(AbstractSource):
    """
    This is a concrete implementation of a Source connector that provides implementations of all the methods needed to run sync
    operations. For simplicity, it also overrides functions that read from files in favor of returning the data directly avoiding
    the need to load static files (ex. spec.yaml, config.json, configured_catalog.json) into the unit-test package.
    """

    def __init__(self, streams: Optional[List[Stream]] = None, authenticator: Optional[AuthBase] = None):
        self._streams = streams
        self._authenticator = authenticator

    def spec(self, logger: logging.Logger) -> ConnectorSpecification:
        return ConnectorSpecification(
            connectionSpecification={
                "$schema": "http://json-schema.org/draft-07/schema#",
                "title": "Test Fixture Spec",
                "type": "object",
                "required": ["api_token"],
                "properties": {
                    "api_token": {
                        "type": "string",
                        "title": "API token",
                        "description": "The token used to authenticate requests to the API.",
                        "airbyte_secret": True,
                    }
                },
            }
        )

    def read_config(self, config_path: str) -> Mapping[str, Any]:
        return {"api_token": "just_some_token"}

    @classmethod
    def read_catalog(cls, catalog_path: str) -> ConfiguredAirbyteCatalog:
        return ConfiguredAirbyteCatalog(
            streams=[
                ConfiguredAirbyteStream(
                    stream=AirbyteStream(
                        name="http_test_stream",
                        json_schema={},
                        supported_sync_modes=[SyncMode.full_refresh, SyncMode.incremental],
                        default_cursor_field=["updated_at"],
                        source_defined_cursor=True,
                        source_defined_primary_key=[["id"]],
                    ),
                    sync_mode=SyncMode.full_refresh,
                    destination_sync_mode=DestinationSyncMode.overwrite,
                )
            ]
        )

    def check_connection(self, *args, **kwargs) -> Tuple[bool, Optional[Any]]:
        return True, ""

    def streams(self, *args, **kwargs) -> List[Stream]:
        return [HttpTestStream(authenticator=self._authenticator)]


class HttpTestStream(HttpStream, ABC):
    url_base = "https://airbyte.com/api/v1/"

    @property
    def cursor_field(self) -> Union[str, List[str]]:
        return ["updated_at"]

    @property
    def availability_strategy(self):
        return None

    def primary_key(self) -> Optional[Union[str, List[str], List[List[str]]]]:
        return "id"

    def path(
        self,
        *,
        stream_state: Mapping[str, Any] = None,
        stream_slice: Mapping[str, Any] = None,
        next_page_token: Mapping[str, Any] = None,
    ) -> str:
        return "cast"

    def parse_response(
        self,
        response: requests.Response,
        *,
        stream_state: Mapping[str, Any],
        stream_slice: Mapping[str, Any] = None,
        next_page_token: Mapping[str, Any] = None,
    ) -> Iterable[Mapping]:
        body = response.json() or {}
        return body["records"]

    def next_page_token(self, response: requests.Response) -> Optional[Mapping[str, Any]]:
        return None

    def get_json_schema(self) -> Mapping[str, Any]:
        return {}


def fixture_mock_send(self, request, **kwargs) -> requests.Response:
    """
    Helper method that can be used by a test to patch the Session.send() function and mock the outbound send operation to provide
    faster and more reliable responses compared to actual API requests
    """
    response = requests.Response()
    response.request = request
    response.status_code = 200
    response.headers = {"header": "value"}
    response_body = {
        "records": [
            {"id": 1, "name": "Celine Song", "position": "director"},
            {"id": 2, "name": "Shabier Kirchner", "position": "cinematographer"},
            {"id": 3, "name": "Christopher Bear", "position": "composer"},
            {"id": 4, "name": "Daniel Rossen", "position": "composer"},
        ]
    }
    response._content = json.dumps(response_body).encode("utf-8")
    return response


class SourceFixtureOauthAuthenticator(Oauth2Authenticator):
    """
    Test OAuth authenticator that only overrides the request and response aspect of the authenticator flow
    """

    def refresh_access_token(self) -> Tuple[str, int]:
        response = requests.request(method="POST", url=self.get_token_refresh_endpoint(), params={})
        response.raise_for_status()
        return "some_access_token", 1800  # Mock oauth response values to be used during the data retrieval step