Live/regression tests: add support for selecting from a subset of connections (#42044)

clnoll · web-flow · commit 548c16cf1e8a · 2024-07-17T11:07:21.000-04:00
diff --git a/.github/workflows/live_tests.yml b/.github/workflows/live_tests.yml
@@ -29,6 +29,13 @@ on:
         description: Use the local CDK when building the target connector
         default: "false"
         type: boolean
+      connection_subset:
+        description: The subset of connections to select from.
+        required: true
+        type: choice
+        options:
+          - sandboxes
+          - all
 
 jobs:
   live_tests:
@@ -86,6 +93,11 @@ jobs:
             echo "USE_LOCAL_CDK_FLAG=" >> $GITHUB_ENV
           fi
 
+      - name: Setup Connection Subset Option
+        if: github.event_name == 'workflow_dispatch'
+        run: |
+          echo "CONNECTION_SUBSET=--connector_live_tests.connection-subset=${{ github.event.inputs.connection_subset }}" >> $GITHUB_ENV
+
       - name: Run Live Tests [WORKFLOW DISPATCH]
         if: github.event_name == 'workflow_dispatch' # TODO: consider using the matrix strategy (https://docs.github.com/en/actions/using-jobs/using-a-matrix-for-your-jobs). See https://github.com/airbytehq/airbyte/pull/37659#discussion_r1583380234 for details.
         uses: ./.github/actions/run-airbyte-ci
@@ -102,4 +114,4 @@ jobs:
           github_token: ${{ secrets.GH_PAT_MAINTENANCE_OSS }}
           s3_build_cache_access_key_id: ${{ secrets.SELF_RUNNER_AWS_ACCESS_KEY_ID }}
           s3_build_cache_secret_key: ${{ secrets.SELF_RUNNER_AWS_SECRET_ACCESS_KEY }}
-          subcommand: connectors ${{ env.USE_LOCAL_CDK_FLAG }} --name ${{ github.event.inputs.connector_name }} test --only-step connector_live_tests --connector_live_tests.test-suite=all --connector_live_tests.connection-id=${{ github.event.inputs.connection_id }} --connector_live_tests.pr-url=${{ github.event.inputs.pr_url }} ${{ env.STREAM_PARAMS }}
+          subcommand: connectors ${{ env.USE_LOCAL_CDK_FLAG }} --name ${{ github.event.inputs.connector_name }} test --only-step connector_live_tests --connector_live_tests.test-suite=all --connector_live_tests.connection-id=${{ github.event.inputs.connection_id }} --connector_live_tests.pr-url=${{ github.event.inputs.pr_url }} ${{ env.STREAM_PARAMS }} ${{ env.CONNECTION_SUBSET }}
diff --git a/.github/workflows/regression_tests.yml b/.github/workflows/regression_tests.yml
@@ -29,6 +29,13 @@ on:
         description: Use the local CDK when building the target connector
         default: "false"
         type: boolean
+      connection_subset:
+        description: The subset of connections to select from.
+        required: true
+        type: choice
+        options:
+          - sandboxes
+          - all
 
 jobs:
   regression_tests:
@@ -86,6 +93,11 @@ jobs:
             echo "USE_LOCAL_CDK_FLAG=" >> $GITHUB_ENV
           fi
 
+      - name: Setup Connection Subset Option
+        if: github.event_name == 'workflow_dispatch'
+        run: |
+          echo "CONNECTION_SUBSET=--connector_live_tests.connection-subset=${{ github.event.inputs.connection_subset }}" >> $GITHUB_ENV
+
       - name: Run Regression Tests [WORKFLOW DISPATCH]
         if: github.event_name == 'workflow_dispatch' # TODO: consider using the matrix strategy (https://docs.github.com/en/actions/using-jobs/using-a-matrix-for-your-jobs). See https://github.com/airbytehq/airbyte/pull/37659#discussion_r1583380234 for details.
         uses: ./.github/actions/run-airbyte-ci
@@ -102,4 +114,4 @@ jobs:
           github_token: ${{ secrets.GH_PAT_MAINTENANCE_OSS }}
           s3_build_cache_access_key_id: ${{ secrets.SELF_RUNNER_AWS_ACCESS_KEY_ID }}
           s3_build_cache_secret_key: ${{ secrets.SELF_RUNNER_AWS_SECRET_ACCESS_KEY }}
-          subcommand: connectors ${{ env.USE_LOCAL_CDK_FLAG }} --name ${{ github.event.inputs.connector_name }} test --only-step connector_live_tests --connector_live_tests.test-suite=regression --connector_live_tests.connection-id=${{ github.event.inputs.connection_id }} --connector_live_tests.pr-url=${{ github.event.inputs.pr_url }} ${{ env.STREAM_PARAMS }}
+          subcommand: connectors ${{ env.USE_LOCAL_CDK_FLAG }} --name ${{ github.event.inputs.connector_name }} test --only-step connector_live_tests --connector_live_tests.test-suite=regression --connector_live_tests.connection-id=${{ github.event.inputs.connection_id }} --connector_live_tests.pr-url=${{ github.event.inputs.pr_url }} ${{ env.STREAM_PARAMS }} ${{ env.CONNECTION_SUBSET }}
diff --git a/airbyte-ci/connectors/live-tests/README.md b/airbyte-ci/connectors/live-tests/README.md
@@ -280,6 +280,10 @@ The traffic recorded on the control connector is passed to the target connector
 
 ## Changelog
 
+### 0.18.0
+
+Add support for selecting from a subset of connections.
+
 ### 0.17.8
 
 Fix the self-signed certificate path we bind to Python connectors.
diff --git a/airbyte-ci/connectors/live-tests/pyproject.toml b/airbyte-ci/connectors/live-tests/pyproject.toml
@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "live-tests"
-version = "0.17.8"
+version = "0.18.0"
 description = "Contains utilities for testing connectors against live data."
 authors = ["Airbyte <contact@airbyte.io>"]
 license = "MIT"
diff --git a/airbyte-ci/connectors/live-tests/src/live_tests/commons/connection_objects_retrieval.py b/airbyte-ci/connectors/live-tests/src/live_tests/commons/connection_objects_retrieval.py
@@ -10,6 +10,7 @@
 import rich
 from connection_retriever import ConnectionObject, retrieve_objects  # type: ignore
 from connection_retriever.errors import NotPermittedError  # type: ignore
+from live_tests.commons.models import ConnectionSubset
 
 from .models import AirbyteCatalog, Command, ConfiguredAirbyteCatalog, ConnectionObjects, SecretDict
 
@@ -95,6 +96,7 @@ def get_connection_objects(
     connector_version: Optional[str] = None,
     auto_select_connection: bool = False,
     selected_streams: Optional[set[str]] = None,
+    connection_subset: ConnectionSubset = ConnectionSubset.SANDBOXES,
 ) -> ConnectionObjects:
     """This function retrieves the connection objects values.
     It checks that the required objects are available and raises a UsageError if they are not.
@@ -113,6 +115,7 @@ def get_connection_objects(
         connector_version (Optional[str]): The version for the connector under test.
         auto_select_connection (bool, optional): Whether to automatically select a connection if no connection id is passed. Defaults to False.
         selected_streams (Optional[Set[str]]): The set of selected streams to use when auto selecting a connection.
+        connection_subset (ConnectionSubset): The subset of connections to select from.
     Raises:
         click.UsageError: If a required object is missing for the command.
         click.UsageError: If a retrieval reason is missing when passing a connection id.
@@ -146,6 +149,7 @@ def get_connection_objects(
             custom_config=custom_config,
             custom_configured_catalog=custom_configured_catalog,
             custom_state=custom_state,
+            connection_subset=connection_subset,
         )
 
     else:
@@ -160,6 +164,7 @@ def get_connection_objects(
                 custom_config=custom_config,
                 custom_configured_catalog=custom_configured_catalog,
                 custom_state=custom_state,
+                connection_subset=connection_subset,
             )
 
         else:
@@ -175,6 +180,7 @@ def get_connection_objects(
                 destination_id=None,
                 connection_id=None,
                 source_docker_image=None,
+                connection_subset=connection_subset,
             )
 
     if fail_if_missing_objects:
@@ -198,6 +204,7 @@ def _get_connection_objects_from_retrieved_objects(
     custom_config: Optional[Dict] = None,
     custom_configured_catalog: Optional[ConfiguredAirbyteCatalog] = None,
     custom_state: Optional[Dict] = None,
+    connection_subset: ConnectionSubset = ConnectionSubset.SANDBOXES,
 ):
     LOGGER.info("Retrieving connection objects from the database...")
     connection_id, retrieved_objects = retrieve_objects(
@@ -208,6 +215,7 @@ def _get_connection_objects_from_retrieved_objects(
         prompt_for_connection_selection=prompt_for_connection_selection,
         with_streams=selected_streams,
         connection_id=connection_id,
+        connection_subset=connection_subset,
     )
 
     retrieved_source_config = parse_config(retrieved_objects.get(ConnectionObject.SOURCE_CONFIG))
diff --git a/airbyte-ci/connectors/live-tests/src/live_tests/commons/models.py b/airbyte-ci/connectors/live-tests/src/live_tests/commons/models.py
@@ -148,6 +148,11 @@ class ActorType(Enum):
     DESTINATION = "destination"
 
 
+class ConnectionSubset(Enum):
+    SANDBOXES = "sandboxes"
+    ALL = "all"
+
+
 @dataclass
 class ConnectorUnderTest:
     image_name: str
diff --git a/airbyte-ci/connectors/live-tests/src/live_tests/conftest.py b/airbyte-ci/connectors/live-tests/src/live_tests/conftest.py
@@ -23,6 +23,7 @@
     ActorType,
     Command,
     ConnectionObjects,
+    ConnectionSubset,
     ConnectorUnderTest,
     ExecutionInputs,
     ExecutionResult,
@@ -90,6 +91,12 @@ def pytest_addoption(parser: Parser) -> None:
         default=TestEvaluationMode.STRICT.value,
         help='If "diagnostic" mode is selected, all tests will pass as long as there is no exception; warnings will be logged. In "strict" mode, tests may fail.',
     )
+    parser.addoption(
+        "--connection-subset",
+        choices=[c.value for c in ConnectionSubset],
+        default=ConnectionSubset.SANDBOXES.value,
+        help="Whether to select from sandbox accounts only.",
+    )
 
 
 def pytest_configure(config: Config) -> None:
@@ -128,6 +135,7 @@ def pytest_configure(config: Config) -> None:
     config.stash[stash_keys.CONNECTOR_IMAGE] = get_option_or_fail(config, "--connector-image")
     config.stash[stash_keys.TARGET_VERSION] = get_option_or_fail(config, "--target-version")
     config.stash[stash_keys.CONTROL_VERSION] = get_control_version(config)
+    config.stash[stash_keys.CONNECTION_SUBSET] = ConnectionSubset(get_option_or_fail(config, "--connection-subset"))
     custom_source_config_path = config.getoption("--config-path")
     custom_configured_catalog_path = config.getoption("--catalog-path")
     custom_state_path = config.getoption("--state-path")
@@ -165,6 +173,7 @@ def pytest_configure(config: Config) -> None:
             connector_version=config.stash[stash_keys.CONTROL_VERSION],
             auto_select_connection=config.stash[stash_keys.AUTO_SELECT_CONNECTION],
             selected_streams=config.stash[stash_keys.SELECTED_STREAMS],
+            connection_subset=config.stash[stash_keys.CONNECTION_SUBSET],
         )
         config.stash[stash_keys.IS_PERMITTED_BOOL] = True
     except (ConnectionNotFoundError, NotPermittedError) as exc:
diff --git a/airbyte-ci/connectors/live-tests/src/live_tests/stash_keys.py b/airbyte-ci/connectors/live-tests/src/live_tests/stash_keys.py
@@ -5,7 +5,7 @@
 
 import pytest
 from live_tests.commons.evaluation_modes import TestEvaluationMode
-from live_tests.commons.models import ConnectionObjects
+from live_tests.commons.models import ConnectionObjects, ConnectionSubset
 from live_tests.report import Report
 
 AIRBYTE_API_KEY = pytest.StashKey[str]()
@@ -15,6 +15,7 @@
 CONNECTION_URL = pytest.StashKey[str | None]()
 CONNECTOR_IMAGE = pytest.StashKey[str]()
 CONTROL_VERSION = pytest.StashKey[str]()
+CONNECTION_SUBSET = pytest.StashKey[ConnectionSubset]()
 DAGGER_LOG_PATH = pytest.StashKey[Path]()
 DUCKDB_PATH = pytest.StashKey[Path]()
 HTTP_DUMP_CACHE_VOLUMES = pytest.StashKey[list]()
diff --git a/airbyte-ci/connectors/pipelines/README.md b/airbyte-ci/connectors/pipelines/README.md
@@ -766,13 +766,14 @@ E.G.: running Poe tasks on the modified internal packages of the current branch:
 ## Changelog
 
 | Version | PR                                                         | Description                                                                                                                  |
-| ------- | ---------------------------------------------------------- | ---------------------------------------------------------------------------------------------------------------------------- |
+|---------| ---------------------------------------------------------- |------------------------------------------------------------------------------------------------------------------------------|
+| 4.25.0  | [#42044](https://github.com/airbytehq/airbyte/pull/42044) | Live/regression tests: add support for selecting from a subset of connections                                                |
 | 4.24.3  | [#42040](https://github.com/airbytehq/airbyte/pull/42040) | Always send regression test approval status check; skip on auto-merge PRs.                                                   |
 | 4.24.2  | [#41676](https://github.com/airbytehq/airbyte/pull/41676) | Send regression test approval status check when skipped.                                                                     |
 | 4.24.1  | [#41642](https://github.com/airbytehq/airbyte/pull/41642) | Use the AIRBYTE_GITHUB_REPO environment variable to run airbyte-ci in other repos.                                           |
 | 4.24.0  | [#41627](https://github.com/airbytehq/airbyte/pull/41627) | Require manual regression test approval for certified connectors                                                             |
 | 4.23.1  | [#41541](https://github.com/airbytehq/airbyte/pull/41541)  | Add support for submodule use-case.                                                                                          |
-| 4.23.0  | [#39906](https://github.com/airbytehq/airbyte/pull/39906)  | Add manifest only build pipeline                                                          |
+| 4.23.0  | [#39906](https://github.com/airbytehq/airbyte/pull/39906)  | Add manifest only build pipeline                                                                                             |
 | 4.22.0  | [#41623](https://github.com/airbytehq/airbyte/pull/41623)  | Make `airbyte-ci` run on private forks.                                                                                      |
 | 4.21.1  | [#41029](https://github.com/airbytehq/airbyte/pull/41029)  | `up-to-date`: mount local docker config to `Syft` to pull private images and benefit from increased DockerHub rate limits.   |
 | 4.21.0  | [#40547](https://github.com/airbytehq/airbyte/pull/40547)  | Make bump-version accept a `--pr-number` option.                                                                             |
diff --git a/airbyte-ci/connectors/pipelines/pipelines/airbyte_ci/connectors/test/steps/common.py b/airbyte-ci/connectors/pipelines/pipelines/airbyte_ci/connectors/test/steps/common.py
@@ -516,6 +516,7 @@ def _get_command_options(self) -> List[str]:
             command_options += ["--test-evaluation-mode", self.test_evaluation_mode]
         if self.selected_streams:
             command_options += ["--stream", self.selected_streams]
+        command_options += ["--connection-subset", self.connection_subset]
         return command_options
 
     def _run_command_with_proxy(self, command: str) -> List[str]:
@@ -565,11 +566,12 @@ def __init__(self, context: ConnectorContext) -> None:
 
         self.test_suite = self.context.run_step_options.get_item_or_default(options, "test-suite", LiveTestSuite.ALL.value)
         self.test_dir = self.test_suite_to_dir[LiveTestSuite(self.test_suite)]
-        self.control_version = self.context.run_step_options.get_item_or_default(options, "control-version", "latest")
+        self.control_version = self.context.run_step_options.get_item_or_default(options, "control-version", None)
         self.target_version = self.context.run_step_options.get_item_or_default(options, "target-version", "dev")
         self.should_read_with_state = self.context.run_step_options.get_item_or_default(options, "should-read-with-state", "1")
         self.selected_streams = self.context.run_step_options.get_item_or_default(options, "selected-streams", None)
         self.test_evaluation_mode = "strict" if self.context.connector.metadata.get("supportLevel") == "certified" else "diagnostic"
+        self.connection_subset = self.context.run_step_options.get_item_or_default(options, "connection-subset", "sandboxes")
         self.run_id = os.getenv("GITHUB_RUN_ID") or str(int(time.time()))
 
     async def _run(self, connector_under_test_container: Container) -> StepResult:
diff --git a/airbyte-ci/connectors/pipelines/pyproject.toml b/airbyte-ci/connectors/pipelines/pyproject.toml
@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "pipelines"
-version = "4.24.3"
+version = "4.25.0"
 description = "Packaged maintained by the connector operations team to perform CI for connectors' pipelines"
 authors = ["Airbyte <contact@airbyte.io>"]
 
