airbyte_ops_mcp.mcp.regression_tests

  1# Copyright (c) 2025 Airbyte, Inc., all rights reserved.
  2"""MCP tools for connector regression tests.
  3
  4This module provides MCP tools for triggering regression tests on Airbyte Cloud
  5connections via GitHub Actions workflows. Regression tests can run in two modes:
  6- Single version mode: Tests a connector version against a connection config
  7- Comparison mode: Compares a target version against a control (baseline) version
  8
  9Tests run asynchronously in GitHub Actions and results can be polled via workflow status.
 10
 11Note: The term "regression tests" encompasses all connector validation testing.
 12The term "live tests" is reserved for scenarios where actual Cloud connections
 13are pinned to pre-release versions for real-world validation.
 14
 15## MCP reference
 16
 17.. include:: ../../../docs/mcp-generated/regression_tests.md
 18    :start-line: 2
 19"""
 20
 21from __future__ import annotations
 22
 23__all__: list[str] = []
 24
 25import uuid
 26from datetime import datetime
 27from enum import Enum
 28from typing import Annotated, Any
 29
 30import requests
 31from airbyte.cloud import CloudWorkspace
 32from airbyte.cloud.auth import resolve_cloud_client_id, resolve_cloud_client_secret
 33from airbyte.exceptions import (
 34    AirbyteMissingResourceError,
 35    AirbyteWorkspaceMismatchError,
 36)
 37from fastmcp import FastMCP
 38from fastmcp_extensions import mcp_tool, register_mcp_tools
 39from pydantic import BaseModel, Field
 40
 41from airbyte_ops_mcp.constants import WorkspaceAliasEnum
 42from airbyte_ops_mcp.github_actions import (
 43    resolve_default_workflow_branch,
 44    trigger_workflow_dispatch,
 45)
 46from airbyte_ops_mcp.github_api import (
 47    GITHUB_API_BASE,
 48    resolve_ci_trigger_github_token,
 49)
 50from airbyte_ops_mcp.mcp.prerelease import ConnectorRepo
 51
 52# =============================================================================
 53# GitHub Workflow Configuration
 54# =============================================================================
 55
 56REGRESSION_TEST_REPO_OWNER = "airbytehq"
 57REGRESSION_TEST_REPO_NAME = "airbyte-ops-mcp"
 58REGRESSION_TEST_DEFAULT_BRANCH = "main"
 59# Unified regression test workflow (handles both single-version and comparison modes)
 60REGRESSION_TEST_WORKFLOW_FILE = "connector-regression-test.yml"
 61
 62
 63# =============================================================================
 64# Workspace Validation Helpers
 65# =============================================================================
 66
 67
 68def validate_connection_workspace(
 69    connection_id: str,
 70    workspace_id: str,
 71) -> None:
 72    """Validate that a connection belongs to the expected workspace.
 73
 74    Uses PyAirbyte's CloudConnection.check_is_valid() method to verify that
 75    the connection exists and belongs to the specified workspace.
 76
 77    Raises:
 78        ValueError: If Airbyte Cloud credentials are missing.
 79        AirbyteWorkspaceMismatchError: If connection belongs to a different workspace.
 80        AirbyteMissingResourceError: If connection is not found.
 81    """
 82    client_id = resolve_cloud_client_id()
 83    client_secret = resolve_cloud_client_secret()
 84    if not client_id or not client_secret:
 85        raise ValueError(
 86            "Missing Airbyte Cloud credentials. "
 87            "Set AIRBYTE_CLOUD_CLIENT_ID and AIRBYTE_CLOUD_CLIENT_SECRET env vars."
 88        )
 89
 90    workspace = CloudWorkspace(
 91        workspace_id=workspace_id,
 92        client_id=client_id,
 93        client_secret=client_secret,
 94    )
 95    connection = workspace.get_connection(connection_id)
 96    connection.check_is_valid()
 97
 98
 99def _get_workflow_run_status(
100    owner: str,
101    repo: str,
102    run_id: int,
103    token: str,
104) -> dict[str, Any]:
105    """Get workflow run details from GitHub API.
106
107    Args:
108        owner: Repository owner (e.g., "airbytehq")
109        repo: Repository name (e.g., "airbyte-ops-mcp")
110        run_id: Workflow run ID
111        token: GitHub API token
112
113    Returns:
114        Workflow run data dictionary.
115
116    Raises:
117        ValueError: If workflow run not found.
118        requests.HTTPError: If API request fails.
119    """
120    url = f"{GITHUB_API_BASE}/repos/{owner}/{repo}/actions/runs/{run_id}"
121    headers = {
122        "Authorization": f"Bearer {token}",
123        "Accept": "application/vnd.github+json",
124        "X-GitHub-Api-Version": "2022-11-28",
125    }
126
127    response = requests.get(url, headers=headers, timeout=30)
128    if response.status_code == 404:
129        raise ValueError(f"Workflow run {owner}/{repo}/actions/runs/{run_id} not found")
130    response.raise_for_status()
131
132    return response.json()
133
134
135# =============================================================================
136# Pydantic Models for Test Results
137# =============================================================================
138
139
140class TestRunStatus(str, Enum):
141    """Status of a test run."""
142
143    QUEUED = "queued"
144    RUNNING = "running"
145    SUCCEEDED = "succeeded"
146    FAILED = "failed"
147
148
149class TestOutcome(str, Enum):
150    """Outcome of a test (execution or comparison)."""
151
152    PENDING = "pending"
153    RUNNING = "running"
154    PASSED = "passed"
155    FAILED = "failed"
156    SKIPPED = "skipped"
157
158
159class ValidationResultModel(BaseModel):
160    """Result of a single validation check."""
161
162    name: str = Field(description="Name of the validation check")
163    passed: bool = Field(description="Whether the validation passed")
164    message: str = Field(description="Human-readable result message")
165    errors: list[str] = Field(
166        default_factory=list,
167        description="List of error messages if validation failed",
168    )
169
170
171class StreamComparisonResultModel(BaseModel):
172    """Result of comparing a single stream between control and target."""
173
174    stream_name: str = Field(description="Name of the stream")
175    passed: bool = Field(description="Whether all comparisons passed")
176    control_record_count: int = Field(description="Number of records in control")
177    target_record_count: int = Field(description="Number of records in target")
178    missing_pks: list[str] = Field(
179        default_factory=list,
180        description="Primary keys present in control but missing in target",
181    )
182    differing_records: int = Field(
183        default=0,
184        description="Number of records that differ between control and target",
185    )
186    message: str = Field(description="Human-readable comparison summary")
187
188
189class RegressionTestExecutionResult(BaseModel):
190    """Results from executing the connector (validations and record counts)."""
191
192    outcome: TestOutcome = Field(description="Outcome of the execution")
193    catalog_validations: list[ValidationResultModel] = Field(
194        default_factory=list,
195        description="Results of catalog validation checks",
196    )
197    record_validations: list[ValidationResultModel] = Field(
198        default_factory=list,
199        description="Results of record validation checks",
200    )
201    record_count: int = Field(
202        default=0,
203        description="Total number of records read",
204    )
205    error_message: str | None = Field(
206        default=None,
207        description="Error message if the execution failed",
208    )
209
210
211class RegressionTestComparisonResult(BaseModel):
212    """Results from comparing target vs control connector versions."""
213
214    outcome: TestOutcome = Field(description="Outcome of the comparison")
215    baseline_version: str | None = Field(
216        default=None,
217        description="Version of the baseline (control) connector",
218    )
219    stream_comparisons: list[StreamComparisonResultModel] = Field(
220        default_factory=list,
221        description="Per-stream comparison results",
222    )
223    error_message: str | None = Field(
224        default=None,
225        description="Error message if the comparison failed",
226    )
227
228
229class RegressionTestResult(BaseModel):
230    """Complete result of a regression test run."""
231
232    run_id: str = Field(description="Unique identifier for this test run")
233    connection_id: str = Field(description="The connection being tested")
234    workspace_id: str = Field(description="The workspace containing the connection")
235    status: TestRunStatus = Field(description="Overall status of the test run")
236    target_version: str | None = Field(
237        default=None,
238        description="Version of the target connector being tested",
239    )
240    baseline_version: str | None = Field(
241        default=None,
242        description="Version of the baseline connector (if comparison mode)",
243    )
244    evaluation_mode: str = Field(
245        default="diagnostic",
246        description="Evaluation mode used (diagnostic or strict)",
247    )
248    compare_versions: bool = Field(
249        default=False,
250        description="Whether comparison mode was used (target vs control)",
251    )
252    execution_result: RegressionTestExecutionResult | None = Field(
253        default=None,
254        description="Results from executing the connector (validations and record counts)",
255    )
256    comparison_result: RegressionTestComparisonResult | None = Field(
257        default=None,
258        description="Results from comparing target vs control connector versions",
259    )
260    artifacts: dict[str, str] = Field(
261        default_factory=dict,
262        description="Paths to generated artifacts (JSONL, DuckDB, HAR files)",
263    )
264    human_summary: str = Field(
265        default="",
266        description="Human-readable summary of the test results",
267    )
268    started_at: datetime | None = Field(
269        default=None,
270        description="When the test run started",
271    )
272    completed_at: datetime | None = Field(
273        default=None,
274        description="When the test run completed",
275    )
276    test_description: str | None = Field(
277        default=None,
278        description="Optional description/context for this test run",
279    )
280
281
282class RunRegressionTestsResponse(BaseModel):
283    """Response from starting a regression test via GitHub Actions workflow."""
284
285    run_id: str = Field(
286        description="Unique identifier for the test run (internal tracking ID)"
287    )
288    status: TestRunStatus = Field(description="Initial status of the test run")
289    message: str = Field(description="Human-readable status message")
290    workflow_url: str | None = Field(
291        default=None,
292        description="URL to view the GitHub Actions workflow file",
293    )
294    github_run_id: int | None = Field(
295        default=None,
296        description="GitHub Actions workflow run ID (use with check_ci_workflow_status)",
297    )
298    github_run_url: str | None = Field(
299        default=None,
300        description="Direct URL to the GitHub Actions workflow run",
301    )
302
303
304# =============================================================================
305# MCP Tools
306# =============================================================================
307
308
309@mcp_tool(
310    read_only=False,
311    idempotent=False,
312    open_world=True,
313)
314def run_regression_tests(
315    connector_name: Annotated[
316        str,
317        "Connector name to build from source (e.g., 'source-pokeapi'). Required.",
318    ],
319    pr: Annotated[
320        int,
321        "PR number to checkout and build from (e.g., 70847). Required. "
322        "The PR must be from the repository specified by the 'repo' parameter.",
323    ],
324    repo: Annotated[
325        ConnectorRepo,
326        "Repository where the connector PR is located. "
327        "Use 'airbyte' for OSS connectors (default) or 'airbyte-enterprise' for enterprise connectors.",
328    ],
329    connection_id: Annotated[
330        str | None,
331        "Airbyte Cloud connection ID to fetch config/catalog from. "
332        "If not provided, uses GSM integration test secrets.",
333    ] = None,
334    skip_compare: Annotated[
335        bool,
336        "If True, skip comparison and run single-version tests only. "
337        "If False (default), run comparison tests (target vs control versions).",
338    ] = False,
339    skip_read_action: Annotated[
340        bool,
341        "If True, skip the read action (run only spec, check, discover). "
342        "If False (default), run all verbs including read.",
343    ] = False,
344    override_test_image: Annotated[
345        str | None,
346        "Override test connector image with tag (e.g., 'airbyte/source-github:1.0.0'). "
347        "Ignored if skip_compare=False.",
348    ] = None,
349    override_control_image: Annotated[
350        str | None,
351        "Override control connector image (baseline version) with tag. "
352        "Ignored if skip_compare=True.",
353    ] = None,
354    workspace_id: Annotated[
355        str | WorkspaceAliasEnum | None,
356        "Optional Airbyte Cloud workspace ID (UUID) or alias. If provided with connection_id, "
357        "validates that the connection belongs to this workspace before triggering tests. "
358        "Accepts '@devin-ai-sandbox' as an alias for the Devin AI sandbox workspace.",
359    ] = None,
360    selected_streams: Annotated[
361        list[str] | None,
362        "List of stream names to include in the read. Only these streams will be included "
363        "in the configured catalog. This is useful to limit data volume by testing only "
364        "specific streams. If not provided, all streams are tested.",
365    ] = None,
366    enable_debug_logs: Annotated[
367        bool,
368        "Enable debug-level logging for regression test output. "
369        "Also passed as `LOG_LEVEL=DEBUG` to the connector Docker container.",
370    ] = False,
371    with_state: Annotated[
372        bool | None,
373        "Fetch and pass the connection's current state to the read command, "
374        "producing a warm read instead of a cold read. Defaults to `True` when "
375        "`connection_id` is provided, `False` otherwise. Has no effect unless "
376        "the command is `read`.",
377    ] = None,
378) -> RunRegressionTestsResponse:
379    """Start a regression test run via GitHub Actions workflow.
380
381    This tool triggers the regression test workflow which builds the connector
382    from the specified PR and runs tests against it.
383
384    Supports both OSS connectors (from airbytehq/airbyte) and enterprise connectors
385    (from airbytehq/airbyte-enterprise). Use the 'repo' parameter to specify which
386    repository contains the connector PR.
387
388    - skip_compare=False (default): Comparison mode - compares the PR version
389      against the baseline (control) version.
390    - skip_compare=True: Single-version mode - runs tests without comparison.
391
392    If connection_id is provided, config/catalog are fetched from Airbyte Cloud.
393    Otherwise, GSM integration test secrets are used.
394
395    Returns immediately with a run_id and workflow URL. Check the workflow URL
396    to monitor progress and view results.
397
398    Requires GITHUB_CI_WORKFLOW_TRIGGER_PAT or GITHUB_TOKEN environment variable
399    with 'actions:write' permission.
400    """
401    # Resolve workspace ID alias
402    resolved_workspace_id = WorkspaceAliasEnum.resolve(workspace_id)
403
404    # Generate a unique run ID for tracking
405    run_id = str(uuid.uuid4())
406
407    # Get GitHub token
408    try:
409        token = resolve_ci_trigger_github_token()
410    except ValueError as e:
411        return RunRegressionTestsResponse(
412            run_id=run_id,
413            status=TestRunStatus.FAILED,
414            message=str(e),
415            workflow_url=None,
416        )
417
418    # Validate workspace membership if workspace_id and connection_id are provided
419    if resolved_workspace_id and connection_id:
420        try:
421            validate_connection_workspace(connection_id, resolved_workspace_id)
422        except (
423            ValueError,
424            AirbyteWorkspaceMismatchError,
425            AirbyteMissingResourceError,
426        ) as e:
427            return RunRegressionTestsResponse(
428                run_id=run_id,
429                status=TestRunStatus.FAILED,
430                message=str(e),
431                workflow_url=None,
432            )
433
434    # Build workflow inputs - connector_name, pr, and repo are required
435    workflow_inputs: dict[str, str] = {
436        "connector_name": connector_name,
437        "pr": str(pr),
438        "repo": repo,
439    }
440
441    # Add optional inputs
442    if connection_id:
443        workflow_inputs["connection_id"] = connection_id
444    if skip_compare:
445        workflow_inputs["skip_compare"] = "true"
446    if skip_read_action:
447        workflow_inputs["skip_read_action"] = "true"
448    if override_test_image:
449        workflow_inputs["override_test_image"] = override_test_image
450    if override_control_image:
451        workflow_inputs["override_control_image"] = override_control_image
452    if selected_streams:
453        workflow_inputs["selected_streams"] = ",".join(selected_streams)
454    if enable_debug_logs:
455        workflow_inputs["enable_debug_logs"] = "true"
456    if with_state is True:
457        workflow_inputs["with_state"] = "true"
458    elif with_state is False:
459        workflow_inputs["with_state"] = "false"
460
461    mode_description = "single-version" if skip_compare else "comparison"
462
463    dispatch_result = trigger_workflow_dispatch(
464        owner=REGRESSION_TEST_REPO_OWNER,
465        repo=REGRESSION_TEST_REPO_NAME,
466        workflow_file=REGRESSION_TEST_WORKFLOW_FILE,
467        ref=resolve_default_workflow_branch(REGRESSION_TEST_DEFAULT_BRANCH),
468        inputs=workflow_inputs,
469        token=token,
470    )
471
472    view_url = dispatch_result.run_url or dispatch_result.workflow_url
473    connection_info = f" for connection {connection_id}" if connection_id else ""
474    repo_info = f" from {repo}" if repo != ConnectorRepo.AIRBYTE else ""
475    return RunRegressionTestsResponse(
476        run_id=run_id,
477        status=TestRunStatus.QUEUED,
478        message=(
479            f"{mode_description.capitalize()} regression test workflow triggered "
480            f"for {connector_name} (PR #{pr}{repo_info}){connection_info}. View progress at: {view_url}"
481        ),
482        workflow_url=dispatch_result.workflow_url,
483        github_run_id=dispatch_result.run_id,
484        github_run_url=dispatch_result.run_url,
485    )
486
487
488# =============================================================================
489# Registration
490# =============================================================================
491
492
493def register_regression_tests_tools(app: FastMCP) -> None:
494    """Register regression tests tools with the FastMCP app.
495
496    Args:
497        app: FastMCP application instance
498    """
499    register_mcp_tools(app, mcp_module=__name__)