From 16000232bbd49481f1e64fee07f73a6825edc8ea Mon Sep 17 00:00:00 2001 From: acaprau <48705707+acaprau@users.noreply.github.com> Date: Fri, 31 Jul 2026 14:46:57 -0700 Subject: [PATCH] feat(connector-stability): credential capability check framework (#13519) --- .../ee/onyx/connectors/capability_checks.py | 89 +++++ .../connectors/capability_checks/models.py | 172 +++++++++ .../connectors/capability_checks/registry.py | 71 ++++ .../connectors/capability_checks/runner.py | 195 ++++++++++ backend/onyx/utils/threadpool_concurrency.py | 4 +- .../connectors/test_ee_capability_checks.py | 148 ++++++++ .../test_capability_check_runner.py | 340 ++++++++++++++++++ .../test_registry_fallbacks.py | 98 +++++ .../test_verdict_aggregation.py | 184 ++++++++++ 9 files changed, 1299 insertions(+), 2 deletions(-) create mode 100644 backend/ee/onyx/connectors/capability_checks.py create mode 100644 backend/onyx/connectors/capability_checks/models.py create mode 100644 backend/onyx/connectors/capability_checks/registry.py create mode 100644 backend/onyx/connectors/capability_checks/runner.py create mode 100644 backend/tests/unit/ee/onyx/connectors/test_ee_capability_checks.py create mode 100644 backend/tests/unit/onyx/connectors/capability_checks/test_capability_check_runner.py create mode 100644 backend/tests/unit/onyx/connectors/capability_checks/test_registry_fallbacks.py create mode 100644 backend/tests/unit/onyx/connectors/capability_checks/test_verdict_aggregation.py diff --git a/backend/ee/onyx/connectors/capability_checks.py b/backend/ee/onyx/connectors/capability_checks.py new file mode 100644 index 0000000000..bb1b143581 --- /dev/null +++ b/backend/ee/onyx/connectors/capability_checks.py @@ -0,0 +1,89 @@ +"""EE implementation of perm-sync capability checks. + +Fetched by ``onyx.connectors.capability_checks.registry`` via +``fetch_ee_implementation_or_noop``. Only the dispatch lives here; the check +implementations stay in the OSS connector modules, mirroring the +``perm_sync_valid.py`` pattern. +""" + +from ee.onyx.external_permissions.sync_params import ( + source_requires_doc_sync, + source_requires_external_group_sync, +) +from onyx.configs.constants import DocumentSource +from onyx.connectors.capability_checks.models import ( + CapabilityCheck, + CapabilityCheckContext, + CredentialCapability, +) + +# Named perm-sync checks per source. Empty at framework stage: per-connector +# work registers named checks here. +_DOC_PERMISSION_SYNC_CHECKS_BY_SOURCE: dict[DocumentSource, list[CapabilityCheck]] = {} + +_EXTERNAL_GROUP_SYNC_CHECKS_BY_SOURCE: dict[DocumentSource, list[CapabilityCheck]] = {} + + +class _PermSyncFallbackCheck(CapabilityCheck): + """ + Baseline check for a perm-sync capability where the source registers no + named checks. + + Wraps the legacy ``validate_perm_sync`` blob, which covers doc sync and + group sync together; the checks for both capabilities are therefore two + instances of this class sharing one check_id, and the runner executes the + probe once per run and mirrors the outcome onto each. + """ + + def __init__( + self, source: DocumentSource, capability: CredentialCapability + ) -> None: + super().__init__( + capability=capability, + check_id=f"{source.value}_perm_sync", + display_name="Permission sync validation", + is_fallback=True, + ) + + def run(self, context: CapabilityCheckContext) -> None: + assert context.connector is not None, ( + "The runner guarantees an instance of a connector." + ) + context.connector.validate_perm_sync() + + +def get_applicable_perm_sync_capabilities( + source: DocumentSource, +) -> set[CredentialCapability]: + """Returns which perm-sync capabilities exist for this source.""" + applicable: set[CredentialCapability] = set() + if source_requires_doc_sync(source): + applicable.add(CredentialCapability.DOC_PERMISSION_SYNC) + if source_requires_external_group_sync(source): + applicable.add(CredentialCapability.EXTERNAL_GROUP_SYNC) + return applicable + + +def get_perm_sync_capability_checks(source: DocumentSource) -> list[CapabilityCheck]: + """Returns the perm-sync capability checks for a source. + + Applicable capabilities with no registered named checks get the shared + ``validate_perm_sync`` fallback so every sync-capable source has day-one + coverage. + """ + applicable = get_applicable_perm_sync_capabilities(source) + registered_by_capability: dict[CredentialCapability, list[CapabilityCheck]] = { + CredentialCapability.DOC_PERMISSION_SYNC: ( + _DOC_PERMISSION_SYNC_CHECKS_BY_SOURCE.get(source, []) + ), + CredentialCapability.EXTERNAL_GROUP_SYNC: ( + _EXTERNAL_GROUP_SYNC_CHECKS_BY_SOURCE.get(source, []) + ), + } + checks: list[CapabilityCheck] = [] + for capability, registered in registered_by_capability.items(): + if registered: + checks.extend(registered) + elif capability in applicable: + checks.append(_PermSyncFallbackCheck(source, capability)) + return checks diff --git a/backend/onyx/connectors/capability_checks/models.py b/backend/onyx/connectors/capability_checks/models.py new file mode 100644 index 0000000000..d291ea2d91 --- /dev/null +++ b/backend/onyx/connectors/capability_checks/models.py @@ -0,0 +1,172 @@ +from abc import ABC, abstractmethod +from collections.abc import Sequence +from enum import Enum +from typing import Any + +from pydantic import BaseModel, ConfigDict + +from onyx.configs.constants import DocumentSource +from onyx.connectors.interfaces import BaseConnector + + +class CredentialCapability(str, Enum): + """ + User-visible capabilities a credential may or may not support for a source. + """ + + INDEXING = "indexing" + DOC_PERMISSION_SYNC = "doc_permission_sync" + EXTERNAL_GROUP_SYNC = "external_group_sync" + + +class CapabilityCheckStatus(str, Enum): + PASSED = "passed" + # A ``ConnectorValidationError``-family failure: real and actionable. + FAILED = "failed" + # Transient or unknown failure (``UnexpectedValidationError`` or any + # unrecognized exception). Mirrors the contract that such errors must never + # be treated as proof the credential is broken. + INDETERMINATE = "indeterminate" + # The check needs state (connector instance or config) that was not + # supplied. + SKIPPED = "skipped" + + +class CapabilityVerdict(str, Enum): + PASSED = "passed" + # Only non-required checks failed: the capability works with caveats. + PASSED_WITH_WARNINGS = "passed_with_warnings" + FAILED = "failed" + INDETERMINATE = "indeterminate" + SKIPPED = "skipped" + # The source does not have this capability (e.g. no group sync for Slack). + NOT_APPLICABLE = "not_applicable" + + +class CapabilityCheckContext(BaseModel): + """Inputs available to a capability check at run time. + + ``connector`` and ``connector_specific_config`` are None for config-less + credential-time runs; the runner skips checks that declare a requirement on + them. + """ + + # ``BaseConnector`` is not a pydantic type; validate by isinstance. + model_config = ConfigDict(arbitrary_types_allowed=True) + + source: DocumentSource + credential_json: dict[str, Any] + connector: BaseConnector | None = None + connector_specific_config: dict[str, Any] | None = None + + +class CapabilityCheck(ABC): + """ + A named probe of one permission/capability assumption a connector makes. + + Concrete checks pass their metadata to ``__init__`` and implement ``run``, + which raises a ``ValidationError``-family exception on failure and returns + nothing on success; the runner maps exceptions to statuses. + """ + + def __init__( + self, + *, + capability: CredentialCapability, + check_id: str, + display_name: str, + required: bool = True, + requires_connector_instance: bool = True, + requires_connector_config: bool = False, + timeout_seconds: float | None = None, + is_fallback: bool = False, + remediation: str | None = None, + docs_link: str | None = None, + ) -> None: + self.capability = capability + self.check_id = check_id + self.display_name = display_name + # Required checks gate the capability verdict; non-required checks only + # downgrade PASSED to PASSED_WITH_WARNINGS (partial capability). + self.required = required + self.requires_connector_instance = requires_connector_instance + self.requires_connector_config = requires_connector_config + # Per-check hang guard in seconds; None falls back to the + # ``CAPABILITY_CHECK_TIMEOUT_SECONDS`` default. Timing out maps to + # INDETERMINATE, never FAILED. + self.timeout_seconds = timeout_seconds + # True for synthesized wrappers around the legacy validation blobs + # (``validate_connector_settings`` / ``validate_perm_sync``) rather than + # named per-permission probes. + self.is_fallback = is_fallback + self.remediation = remediation + self.docs_link = docs_link + + @abstractmethod + def run(self, context: CapabilityCheckContext) -> None: + """Probes the capability; raises on failure, returns on success.""" + + +class CapabilityCheckResult(BaseModel): + capability: CredentialCapability + check_id: str + display_name: str + required: bool + status: CapabilityCheckStatus + # Exception text for failures, skip reason for skips, empty on success. + message: str = "" + # Exception class name, for support diagnosis. + error_type: str | None = None + is_fallback: bool = False + remediation: str | None = None + docs_link: str | None = None + duration_ms: int | None = None + + +def aggregate_capability_verdict( + applicable: bool, + results: Sequence[CapabilityCheckResult], +) -> CapabilityVerdict: + """Rolls the per-check results of one capability up into a single verdict. + + Pure function. Required checks gate the verdict: a required failure is + FAILED, and a required indeterminate is INDETERMINATE, since the + capability's core is unverified and no PASSED claim can be made. + Non-required outcomes only downgrade: failures to PASSED_WITH_WARNINGS + (definite, stable, actionable), indeterminates to INDETERMINATE when no + warning outranks them (transient, self-resolves on re-run). An empty result + list aggregates to SKIPPED. + """ + if not applicable: + return CapabilityVerdict.NOT_APPLICABLE + if any( + result.status == CapabilityCheckStatus.FAILED and result.required + for result in results + ): + return CapabilityVerdict.FAILED + if any( + result.status == CapabilityCheckStatus.INDETERMINATE and result.required + for result in results + ): + return CapabilityVerdict.INDETERMINATE + if any(result.status == CapabilityCheckStatus.FAILED for result in results): + return CapabilityVerdict.PASSED_WITH_WARNINGS + if any(result.status == CapabilityCheckStatus.INDETERMINATE for result in results): + return CapabilityVerdict.INDETERMINATE + if all(result.status == CapabilityCheckStatus.SKIPPED for result in results): + return CapabilityVerdict.SKIPPED + return CapabilityVerdict.PASSED + + +def compute_capability_verdicts( + applicable_capabilities: set[CredentialCapability], + results: Sequence[CapabilityCheckResult], +) -> dict[CredentialCapability, CapabilityVerdict]: + """Computes a verdict for every capability, applicable or not.""" + return { + capability: aggregate_capability_verdict( + capability in applicable_capabilities, + [result for result in results if result.capability == capability], + ) + for capability in CredentialCapability + } diff --git a/backend/onyx/connectors/capability_checks/registry.py b/backend/onyx/connectors/capability_checks/registry.py new file mode 100644 index 0000000000..1ea7d964bf --- /dev/null +++ b/backend/onyx/connectors/capability_checks/registry.py @@ -0,0 +1,71 @@ +from onyx.configs.constants import DocumentSource +from onyx.connectors.capability_checks.models import ( + CapabilityCheck, + CapabilityCheckContext, + CredentialCapability, +) +from onyx.utils.variable_functionality import fetch_ee_implementation_or_noop + +# INDEXING checks per source. Checks must be enumerable without an instantiated +# connector, hence a registry module rather than a ``BaseConnector`` method. +# Empty at framework stage: per-connector work registers named checks here. +_INDEXING_CHECKS_BY_SOURCE: dict[DocumentSource, list[CapabilityCheck]] = {} + + +class _ConnectorSettingsFallbackCheck(CapabilityCheck): + """Baseline INDEXING check where the source registers no named checks. + + Wraps the legacy ``validate_connector_settings`` blob so every connector has + day-one coverage until a per-connector session registers named, + per-permission checks that shadow this. + """ + + def __init__(self, source: DocumentSource) -> None: + super().__init__( + capability=CredentialCapability.INDEXING, + check_id=f"{source.value}_connector_settings", + display_name="Connector settings validation", + is_fallback=True, + ) + + def run(self, context: CapabilityCheckContext) -> None: + assert context.connector is not None, ( + "The runner guarantees an instance of a connector." + ) + context.connector.validate_connector_settings() + + +def get_capability_checks(source: DocumentSource) -> list[CapabilityCheck]: + """ + Returns all capability checks for a source, synthesizing fallbacks where + there are no named checks. + + INDEXING checks are registered here; perm-sync checks come from the EE + implementation and are empty on OSS builds, where the perm-sync feature does + not exist. + """ + checks = list(_INDEXING_CHECKS_BY_SOURCE.get(source, [])) + if not checks: + checks.append(_ConnectorSettingsFallbackCheck(source)) + get_perm_sync_checks = fetch_ee_implementation_or_noop( + "onyx.connectors.capability_checks", + "get_perm_sync_capability_checks", + noop_return_value=[], + ) + checks.extend(get_perm_sync_checks(source)) + return checks + + +def get_applicable_capabilities(source: DocumentSource) -> set[CredentialCapability]: + """Returns the capabilities that exist for this source on this build. + + Capabilities outside this set aggregate to a NOT_APPLICABLE verdict. On OSS + builds only INDEXING applies, which is correct since perm sync does not + exist there. + """ + get_perm_sync_capabilities = fetch_ee_implementation_or_noop( + "onyx.connectors.capability_checks", + "get_applicable_perm_sync_capabilities", + noop_return_value=set(), + ) + return {CredentialCapability.INDEXING} | get_perm_sync_capabilities(source) diff --git a/backend/onyx/connectors/capability_checks/runner.py b/backend/onyx/connectors/capability_checks/runner.py new file mode 100644 index 0000000000..9e8e94280a --- /dev/null +++ b/backend/onyx/connectors/capability_checks/runner.py @@ -0,0 +1,195 @@ +import time +from collections.abc import Sequence + +from pydantic import BaseModel + +from onyx.connectors.capability_checks.models import ( + CapabilityCheck, + CapabilityCheckContext, + CapabilityCheckResult, + CapabilityCheckStatus, + CredentialCapability, +) +from onyx.connectors.exceptions import ( + ConnectorValidationError, + UnexpectedValidationError, +) +from onyx.utils.logger import setup_logger +from onyx.utils.threadpool_concurrency import run_with_timeout + +logger = setup_logger() + +# Default per-check hang guard. Deliberately long: checks run async with no +# total budget, and a slow probe is vastly cheaper than a failed indexing run. +# Its sole purpose is to stop a stuck probe from blocking a worker thread +# forever. Known-slow probes override it via +# ``CapabilityCheck.timeout_seconds``. +CAPABILITY_CHECK_TIMEOUT_SECONDS = 600 + +_SKIP_NEEDS_INSTANCE_MESSAGE = ( + "Requires a connector instance -- will re-run automatically when the " + "connector is instantiated." +) +_SKIP_NEEDS_CONFIG_MESSAGE = ( + "Requires connector configuration -- will re-run automatically when the " + "connector is configured." +) +_TIMEOUT_MESSAGE = ( + "Check timed out; the source may be slow -- try re-running the checks." +) + + +class _CheckOutcome(BaseModel): + """ + Outcome of one check execution, before it is joined with the check's + metadata into a result row. Shared across checks with one check_id. + """ + + status: CapabilityCheckStatus + message: str = "" + error_type: str | None = None + duration_ms: int | None = None + + +def _build_result( + check: CapabilityCheck, outcome: _CheckOutcome +) -> CapabilityCheckResult: + return CapabilityCheckResult( + capability=check.capability, + check_id=check.check_id, + display_name=check.display_name, + required=check.required, + status=outcome.status, + message=outcome.message, + error_type=outcome.error_type, + is_fallback=check.is_fallback, + remediation=check.remediation, + docs_link=check.docs_link, + duration_ms=outcome.duration_ms, + ) + + +def _execute_check( + check: CapabilityCheck, context: CapabilityCheckContext +) -> _CheckOutcome: + """ + Executes one check under its hang guard and maps the outcome to a status. + + A timeout maps to INDETERMINATE, never FAILED; a slow source is not proof of + a broken credential. + """ + timeout_seconds = check.timeout_seconds or CAPABILITY_CHECK_TIMEOUT_SECONDS + start = time.monotonic() + + def elapsed_ms() -> int: + return int((time.monotonic() - start) * 1000) + + try: + run_with_timeout(timeout_seconds, check.run, context) + except ConnectorValidationError as e: + return _CheckOutcome( + status=CapabilityCheckStatus.FAILED, + message=str(e), + error_type=type(e).__name__, + duration_ms=elapsed_ms(), + ) + except TimeoutError as e: + # ``run_with_timeout`` abandons the probe thread on timeout -- Python + # has no safe thread cancellation -- so the probe may still have a + # request in flight. Check bodies must keep per-request timeouts below + # ``timeout_seconds`` so an abandoned probe dies with its current + # request instead of running on. + logger.warning( + "Capability check %s timed out after %.0fs; its probe thread is " + "abandoned and may still have a request in-flight.", + check.check_id, + timeout_seconds, + ) + return _CheckOutcome( + status=CapabilityCheckStatus.INDETERMINATE, + message=_TIMEOUT_MESSAGE, + error_type=type(e).__name__, + duration_ms=elapsed_ms(), + ) + except UnexpectedValidationError as e: + return _CheckOutcome( + status=CapabilityCheckStatus.INDETERMINATE, + message=str(e), + error_type=type(e).__name__, + duration_ms=elapsed_ms(), + ) + except Exception as e: + logger.warning( + "Capability check %s raised an unexpected error: %s", + check.check_id, + e, + ) + return _CheckOutcome( + status=CapabilityCheckStatus.INDETERMINATE, + message=str(e), + error_type=type(e).__name__, + duration_ms=elapsed_ms(), + ) + return _CheckOutcome(status=CapabilityCheckStatus.PASSED, duration_ms=elapsed_ms()) + + +def run_capability_checks( + checks: Sequence[CapabilityCheck], + context: CapabilityCheckContext, +) -> list[CapabilityCheckResult]: + """Runs checks sequentially and maps their outcomes to statuses. + + Sequential on purpose: source APIs can rate limit aggressively, so + concurrent probes against the same credential are counterproductive. Check + failures become result rows; this function never raises for them. + + Checks sharing one check_id (the perm-sync fallback registered under both + sync capabilities) are one check surfaced per capability: they execute once + and the outcome mirrors onto each result. Distinct check_ids always execute + independently, even when they share an implementation. + """ + # A check_id may repeat only as one check mirrored across capabilities; + # anything else is a registration bug, caught before it silently collapses + # distinct checks into one execution. + check_class_by_id: dict[str, type[CapabilityCheck]] = {} + seen_capability_ids: set[tuple[str, CredentialCapability]] = set() + for check in checks: + capability_id = (check.check_id, check.capability) + assert capability_id not in seen_capability_ids, ( + f"Duplicate check_id {check.check_id!r} within capability " + f"{check.capability.value}." + ) + seen_capability_ids.add(capability_id) + check_class = check_class_by_id.setdefault(check.check_id, type(check)) + assert check_class is type(check), ( + f"check_id {check.check_id!r} is shared by {check_class.__name__} " + f"and {type(check).__name__}; mirrored checks must be one class." + ) + + results: list[CapabilityCheckResult] = [] + outcome_by_check_id: dict[str, _CheckOutcome] = {} + for check in checks: + skip_message: str | None = None + if ( + check.requires_connector_config + and context.connector_specific_config is None + ): + skip_message = _SKIP_NEEDS_CONFIG_MESSAGE + elif check.requires_connector_instance and context.connector is None: + skip_message = _SKIP_NEEDS_INSTANCE_MESSAGE + if skip_message is not None: + results.append( + _build_result( + check, + _CheckOutcome( + status=CapabilityCheckStatus.SKIPPED, + message=skip_message, + ), + ) + ) + continue + + if check.check_id not in outcome_by_check_id: + outcome_by_check_id[check.check_id] = _execute_check(check, context) + results.append(_build_result(check, outcome_by_check_id[check.check_id])) + return results diff --git a/backend/onyx/utils/threadpool_concurrency.py b/backend/onyx/utils/threadpool_concurrency.py index 1761408eca..445215bc17 100644 --- a/backend/onyx/utils/threadpool_concurrency.py +++ b/backend/onyx/utils/threadpool_concurrency.py @@ -526,8 +526,8 @@ def run_with_timeout( timeout: float, func: Callable[..., R], *args: Any, **kwargs: Any ) -> R: """ - Executes a function with a timeout. If the function doesn't complete within the specified - timeout, raises TimeoutError. + Executes a function with a timeout. If the function doesn't complete within + the specified timeout, raises TimeoutError. """ context = contextvars.copy_context() task = TimeoutThread(timeout, context.run, func, *args, **kwargs) diff --git a/backend/tests/unit/ee/onyx/connectors/test_ee_capability_checks.py b/backend/tests/unit/ee/onyx/connectors/test_ee_capability_checks.py new file mode 100644 index 0000000000..1bd45441c0 --- /dev/null +++ b/backend/tests/unit/ee/onyx/connectors/test_ee_capability_checks.py @@ -0,0 +1,148 @@ +from unittest.mock import MagicMock + +import pytest + +from ee.onyx.connectors import capability_checks as ee_capability_checks +from ee.onyx.connectors.capability_checks import ( + get_applicable_perm_sync_capabilities, + get_perm_sync_capability_checks, +) +from onyx.configs.constants import DocumentSource +from onyx.connectors.capability_checks.models import ( + CapabilityCheck, + CapabilityCheckContext, + CredentialCapability, +) +from onyx.connectors.interfaces import BaseConnector + + +class _NamedCheck(CapabilityCheck): + """Minimal concrete named check; ``run`` never executes in these tests.""" + + def run(self, context: CapabilityCheckContext) -> None: + raise NotImplementedError + + +def test_slack_perm_sync_capabilities_exclude_group_sync() -> None: + """Verifies Slack applicability: doc sync yes, group sync no (by design).""" + # Under test and postcondition. + assert get_applicable_perm_sync_capabilities(DocumentSource.SLACK) == { + CredentialCapability.DOC_PERMISSION_SYNC + } + + +def test_google_drive_has_both_perm_sync_capabilities() -> None: + """Verifies a both-capability source resolves from the sync registry.""" + # Under test and postcondition. + assert get_applicable_perm_sync_capabilities(DocumentSource.GOOGLE_DRIVE) == { + CredentialCapability.DOC_PERMISSION_SYNC, + CredentialCapability.EXTERNAL_GROUP_SYNC, + } + + +def test_censoring_only_source_has_no_perm_sync_capabilities() -> None: + """ + Verifies Salesforce (query-time censoring only) maps to no capabilities. + """ + # Under test and postcondition. + assert get_applicable_perm_sync_capabilities(DocumentSource.SALESFORCE) == set() + + +def test_fallback_synthesis_respects_applicability() -> None: + """ + Verifies only applicable capabilities get a fallback (Slack has no group + sync by design, so only the doc-sync fallback is synthesized). + """ + # Under test. + checks = get_perm_sync_capability_checks(DocumentSource.SLACK) + + # Postcondition. + assert [check.capability for check in checks] == [ + CredentialCapability.DOC_PERMISSION_SYNC + ] + assert checks[0].is_fallback is True + assert checks[0].check_id == "slack_perm_sync" + + +def test_unregistered_sync_source_gets_shared_fallback_checks() -> None: + """ + Verifies fallback synthesis for a sync-capable source with no named checks. + """ + # Under test. + checks = get_perm_sync_capability_checks(DocumentSource.GOOGLE_DRIVE) + + # Postcondition. + # One fallback per applicable capability, sharing a single check_id (and + # class) so the runner executes it once and mirrors. + assert {check.capability for check in checks} == { + CredentialCapability.DOC_PERMISSION_SYNC, + CredentialCapability.EXTERNAL_GROUP_SYNC, + } + assert all(check.is_fallback for check in checks) + assert len({type(check) for check in checks}) == 1 + assert {check.check_id for check in checks} == {"google_drive_perm_sync"} + + +def test_registered_checks_clobber_only_their_capability( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """ + Verifies named checks clobber the fallback per capability, not per source. + """ + # Precondition. + # Nothing is registered at framework stage, so register a named doc-sync + # check the way a per-connector session would. + named_check = _NamedCheck( + capability=CredentialCapability.DOC_PERMISSION_SYNC, + check_id="google_drive_named_check", + display_name="Named check", + ) + monkeypatch.setitem( + ee_capability_checks._DOC_PERMISSION_SYNC_CHECKS_BY_SOURCE, + DocumentSource.GOOGLE_DRIVE, + [named_check], + ) + + # Under test. + checks = get_perm_sync_capability_checks(DocumentSource.GOOGLE_DRIVE) + + # Postcondition. + # Doc sync uses the named check; group sync, applicable but still + # unregistered, keeps its fallback. + doc_sync_checks = [ + check + for check in checks + if check.capability == CredentialCapability.DOC_PERMISSION_SYNC + ] + group_sync_checks = [ + check + for check in checks + if check.capability == CredentialCapability.EXTERNAL_GROUP_SYNC + ] + assert doc_sync_checks == [named_check] + assert len(group_sync_checks) == 1 + assert group_sync_checks[0].is_fallback is True + + +def test_censoring_only_source_gets_no_fallback_checks() -> None: + """Verifies no fallback is synthesized when neither capability applies.""" + # Under test and postcondition. + assert get_perm_sync_capability_checks(DocumentSource.SALESFORCE) == [] + + +def test_fallback_check_calls_validate_perm_sync() -> None: + """Verifies the synthesized fallback wraps ``validate_perm_sync``.""" + # Precondition. + connector = MagicMock(spec=BaseConnector) + context = CapabilityCheckContext( + source=DocumentSource.GOOGLE_DRIVE, + credential_json={}, + connector=connector, + ) + check = get_perm_sync_capability_checks(DocumentSource.GOOGLE_DRIVE)[0] + + # Under test. + check.run(context) + + # Postcondition. + connector.validate_perm_sync.assert_called_once_with() diff --git a/backend/tests/unit/onyx/connectors/capability_checks/test_capability_check_runner.py b/backend/tests/unit/onyx/connectors/capability_checks/test_capability_check_runner.py new file mode 100644 index 0000000000..7871efc8a2 --- /dev/null +++ b/backend/tests/unit/onyx/connectors/capability_checks/test_capability_check_runner.py @@ -0,0 +1,340 @@ +import time +from collections.abc import Callable +from typing import cast +from unittest.mock import MagicMock + +import pytest + +from onyx.configs.constants import DocumentSource +from onyx.connectors.capability_checks.models import ( + CapabilityCheck, + CapabilityCheckContext, + CapabilityCheckStatus, + CredentialCapability, +) +from onyx.connectors.capability_checks.runner import run_capability_checks +from onyx.connectors.exceptions import ( + InsufficientPermissionsError, + UnexpectedValidationError, +) +from onyx.connectors.interfaces import BaseConnector + + +class _CallableCheck(CapabilityCheck): + """Concrete check that delegates ``run`` to an injected callable.""" + + def __init__( + self, + run: Callable[[CapabilityCheckContext], None], + check_id: str = "dummy_check", + required: bool = True, + requires_connector_instance: bool = False, + requires_connector_config: bool = False, + capability: CredentialCapability = CredentialCapability.INDEXING, + timeout_seconds: float | None = None, + is_fallback: bool = False, + ) -> None: + super().__init__( + capability=capability, + check_id=check_id, + display_name="Dummy check", + required=required, + requires_connector_instance=requires_connector_instance, + requires_connector_config=requires_connector_config, + timeout_seconds=timeout_seconds, + is_fallback=is_fallback, + ) + self._run = run + + def run(self, context: CapabilityCheckContext) -> None: + self._run(context) + + +class _OtherCallableCheck(_CallableCheck): + """Second concrete class, for the mirrored-check class-identity guard.""" + + +def _context( + connector: BaseConnector | None = None, + connector_specific_config: dict | None = None, +) -> CapabilityCheckContext: + return CapabilityCheckContext( + source=DocumentSource.SLACK, + credential_json={}, + connector=connector, + connector_specific_config=connector_specific_config, + ) + + +def test_passing_check_maps_to_passed() -> None: + """Verifies that a check returning without raising is PASSED.""" + # Precondition. + check = _CallableCheck(lambda _context: None) + + # Under test. + results = run_capability_checks([check], _context()) + + # Postcondition. + assert len(results) == 1 + assert results[0].status == CapabilityCheckStatus.PASSED + assert results[0].message == "" + assert results[0].error_type is None + assert results[0].duration_ms is not None + + +def test_connector_validation_error_maps_to_failed() -> None: + """Verifies that the ``ConnectorValidationError`` family is FAILED.""" + + # Precondition. + def run(_context: CapabilityCheckContext) -> None: + raise InsufficientPermissionsError("Missing `channels:read`.") + + check = _CallableCheck(run) + + # Under test. + results = run_capability_checks([check], _context()) + + # Postcondition. + assert results[0].status == CapabilityCheckStatus.FAILED + assert results[0].error_type == "InsufficientPermissionsError" + assert "channels:read" in results[0].message + + +def test_unexpected_validation_error_maps_to_indeterminate() -> None: + """Verifies that ``UnexpectedValidationError`` is transient, not FAILED.""" + + # Precondition. + def run(_context: CapabilityCheckContext) -> None: + raise UnexpectedValidationError("Slack rate limited the check.") + + check = _CallableCheck(run) + + # Under test. + results = run_capability_checks([check], _context()) + + # Postcondition. + assert results[0].status == CapabilityCheckStatus.INDETERMINATE + assert results[0].error_type == "UnexpectedValidationError" + + +def test_unknown_exception_maps_to_indeterminate() -> None: + """Verifies that unrecognized exceptions never count as real failures.""" + + # Precondition. + def run(_context: CapabilityCheckContext) -> None: + raise RuntimeError("Boom.") + + check = _CallableCheck(run) + + # Under test. + results = run_capability_checks([check], _context()) + + # Postcondition. + assert results[0].status == CapabilityCheckStatus.INDETERMINATE + assert results[0].error_type == "RuntimeError" + + +def test_skips_instance_requiring_check_without_connector() -> None: + """Verifies the skip path for checks that need a connector instance.""" + # Precondition. + run = MagicMock() + check = _CallableCheck(run, requires_connector_instance=True) + + # Under test. + results = run_capability_checks([check], _context(connector=None)) + + # Postcondition. + assert results[0].status == CapabilityCheckStatus.SKIPPED + run.assert_not_called() + + +def test_skips_config_requiring_check_even_with_connector() -> None: + """ + Verifies that config-requiring checks skip when only an instance exists. + """ + # Precondition. + run = MagicMock() + check = _CallableCheck(run, requires_connector_config=True) + connector = cast(BaseConnector, MagicMock(spec=BaseConnector)) + + # Under test. + results = run_capability_checks( + [check], _context(connector=connector, connector_specific_config=None) + ) + + # Postcondition. + assert results[0].status == CapabilityCheckStatus.SKIPPED + run.assert_not_called() + + +def test_config_requiring_check_runs_with_config() -> None: + """Verifies that a supplied config unlocks config-requiring checks.""" + # Precondition. + run = MagicMock(return_value=None) + check = _CallableCheck(run, requires_connector_config=True) + + # Under test. + results = run_capability_checks( + [check], _context(connector_specific_config={"channels": ["general"]}) + ) + + # Postcondition. + assert results[0].status == CapabilityCheckStatus.PASSED + run.assert_called_once() + + +def test_failure_does_not_stop_subsequent_checks() -> None: + """Verifies that the runner continues past failing checks.""" + + # Precondition. + def failing_run(_context: CapabilityCheckContext) -> None: + raise InsufficientPermissionsError("Missing scope.") + + checks = [ + _CallableCheck(failing_run, check_id="failing_check"), + _CallableCheck(lambda _context: None, check_id="passing_check"), + ] + + # Under test. + results = run_capability_checks(checks, _context()) + + # Postcondition. + assert [result.status for result in results] == [ + CapabilityCheckStatus.FAILED, + CapabilityCheckStatus.PASSED, + ] + assert [result.check_id for result in results] == [ + "failing_check", + "passing_check", + ] + + +def test_timeout_maps_to_indeterminate() -> None: + """ + Verifies that the per-check hang guard yields INDETERMINATE, not FAILED. + """ + # Precondition. + # The sleep must exceed the timeout. + timeout_seconds = 0.5 + slow_run_duration_seconds = 1 + assert slow_run_duration_seconds > timeout_seconds + + def slow_run(_context: CapabilityCheckContext) -> None: + time.sleep(slow_run_duration_seconds) + + check = _CallableCheck(slow_run, timeout_seconds=timeout_seconds) + + # Under test. + results = run_capability_checks([check], _context()) + + # Postcondition. + assert results[0].status == CapabilityCheckStatus.INDETERMINATE + assert results[0].error_type == "TimeoutError" + assert "timed out" in results[0].message + + +def test_shared_check_id_executes_once_and_mirrors() -> None: + """ + Verifies the shared perm-sync fallback contract: one execution, two results. + """ + # Precondition. + # One check_id registered under both sync capabilities. + call_count = 0 + + def shared_run(_context: CapabilityCheckContext) -> None: + nonlocal call_count + call_count += 1 + raise InsufficientPermissionsError("Missing directory scope.") + + checks = [ + _CallableCheck( + shared_run, + check_id="perm_sync", + capability=CredentialCapability.DOC_PERMISSION_SYNC, + ), + _CallableCheck( + shared_run, + check_id="perm_sync", + capability=CredentialCapability.EXTERNAL_GROUP_SYNC, + ), + ] + + # Under test. + results = run_capability_checks(checks, _context()) + + # Postcondition. + # Both capabilities carry the mirrored outcome. + assert call_count == 1 + assert [result.status for result in results] == [ + CapabilityCheckStatus.FAILED, + CapabilityCheckStatus.FAILED, + ] + assert {result.capability for result in results} == { + CredentialCapability.DOC_PERMISSION_SYNC, + CredentialCapability.EXTERNAL_GROUP_SYNC, + } + assert all("directory scope" in result.message for result in results) + + +def test_shared_callable_with_distinct_check_ids_runs_independently() -> None: + """ + Verifies distinct checks are never collapsed by a shared implementation. + """ + # Precondition. + shared_run = MagicMock(return_value=None) + checks = [ + _CallableCheck(shared_run, check_id="first_check"), + _CallableCheck(shared_run, check_id="second_check"), + ] + + # Under test. + run_capability_checks(checks, _context()) + + # Postcondition. + assert shared_run.call_count == 2 + + +def test_duplicate_check_id_within_a_capability_is_rejected() -> None: + """Verifies a reused check_id within one capability is a hard error.""" + # Precondition. + checks = [ + _CallableCheck(MagicMock(return_value=None), check_id="same_id"), + _CallableCheck(MagicMock(return_value=None), check_id="same_id"), + ] + + # Under test and postcondition. + with pytest.raises(AssertionError, match="Duplicate check_id"): + run_capability_checks(checks, _context()) + + +def test_mirrored_check_id_across_classes_is_rejected() -> None: + """Verifies cross-capability check_id reuse requires a single class.""" + # Precondition. + checks = [ + _CallableCheck( + MagicMock(return_value=None), + check_id="perm_sync", + capability=CredentialCapability.DOC_PERMISSION_SYNC, + ), + _OtherCallableCheck( + MagicMock(return_value=None), + check_id="perm_sync", + capability=CredentialCapability.EXTERNAL_GROUP_SYNC, + ), + ] + + # Under test and postcondition. + with pytest.raises(AssertionError, match="must be one class"): + run_capability_checks(checks, _context()) + + +def test_is_fallback_propagates_to_result() -> None: + """Verifies that result rows carry the fallback marker for FE labeling.""" + # Precondition. + check = _CallableCheck(lambda _context: None, is_fallback=True) + + # Under test. + results = run_capability_checks([check], _context()) + + # Postcondition. + assert results[0].is_fallback is True diff --git a/backend/tests/unit/onyx/connectors/capability_checks/test_registry_fallbacks.py b/backend/tests/unit/onyx/connectors/capability_checks/test_registry_fallbacks.py new file mode 100644 index 0000000000..0d1ad64fe5 --- /dev/null +++ b/backend/tests/unit/onyx/connectors/capability_checks/test_registry_fallbacks.py @@ -0,0 +1,98 @@ +from unittest.mock import MagicMock + +import pytest + +from onyx.configs.constants import DocumentSource +from onyx.connectors.capability_checks import registry +from onyx.connectors.capability_checks.models import ( + CapabilityCheck, + CapabilityCheckContext, + CredentialCapability, +) +from onyx.connectors.capability_checks.registry import ( + get_applicable_capabilities, + get_capability_checks, +) +from onyx.connectors.interfaces import BaseConnector + + +class _NamedCheck(CapabilityCheck): + """Minimal concrete named check; ``run`` never executes in these tests.""" + + def run(self, context: CapabilityCheckContext) -> None: + raise NotImplementedError + + +def test_unregistered_source_gets_connector_settings_fallback() -> None: + """ + Verifies fallback synthesis for a source with no named INDEXING checks. + """ + # Under test. + checks = get_capability_checks(DocumentSource.GITHUB) + + # Postcondition. + # On OSS builds only the synthesized INDEXING check exists. + indexing_checks = [ + check for check in checks if check.capability == CredentialCapability.INDEXING + ] + assert len(indexing_checks) == 1 + fallback = indexing_checks[0] + assert fallback.check_id == "github_connector_settings" + assert fallback.is_fallback is True + assert fallback.required is True + assert fallback.requires_connector_instance is True + + +def test_registered_source_gets_no_fallback(monkeypatch: pytest.MonkeyPatch) -> None: + """Verifies that registered named checks clobber the INDEXING fallback.""" + # Precondition. + # Nothing is registered at framework stage, so register a named check the + # way a per-connector session would. + named_check = _NamedCheck( + capability=CredentialCapability.INDEXING, + check_id="github_named_check", + display_name="Named check", + ) + monkeypatch.setitem( + registry._INDEXING_CHECKS_BY_SOURCE, DocumentSource.GITHUB, [named_check] + ) + + # Under test. + checks = get_capability_checks(DocumentSource.GITHUB) + + # Postcondition. + indexing_checks = [ + check for check in checks if check.capability == CredentialCapability.INDEXING + ] + assert indexing_checks == [named_check] + assert not any(check.is_fallback for check in indexing_checks) + + +def test_fallback_check_calls_validate_connector_settings() -> None: + """ + Verifies the synthesized fallback wraps ``validate_connector_settings``. + """ + # Precondition. + connector = MagicMock(spec=BaseConnector) + context = CapabilityCheckContext( + source=DocumentSource.GITHUB, + credential_json={}, + connector=connector, + ) + check = get_capability_checks(DocumentSource.GITHUB)[0] + + # Under test. + check.run(context) + + # Postcondition. + connector.validate_connector_settings.assert_called_once_with() + + +def test_oss_build_only_has_indexing_capability() -> None: + """Verifies the EE-noop path: perm-sync capabilities do not exist on OSS.""" + # Under test and postcondition. + # Google Drive is sync-capable on EE, so it is the strongest witness that + # the OSS noop returns INDEXING only. + assert get_applicable_capabilities(DocumentSource.GOOGLE_DRIVE) == { + CredentialCapability.INDEXING + } diff --git a/backend/tests/unit/onyx/connectors/capability_checks/test_verdict_aggregation.py b/backend/tests/unit/onyx/connectors/capability_checks/test_verdict_aggregation.py new file mode 100644 index 0000000000..68ddc735d7 --- /dev/null +++ b/backend/tests/unit/onyx/connectors/capability_checks/test_verdict_aggregation.py @@ -0,0 +1,184 @@ +from onyx.connectors.capability_checks.models import ( + CapabilityCheckResult, + CapabilityCheckStatus, + CapabilityVerdict, + CredentialCapability, + aggregate_capability_verdict, + compute_capability_verdicts, +) + + +def _result( + status: CapabilityCheckStatus, + required: bool = True, + capability: CredentialCapability = CredentialCapability.INDEXING, +) -> CapabilityCheckResult: + return CapabilityCheckResult( + capability=capability, + check_id="check", + display_name="Check", + required=required, + status=status, + ) + + +def test_not_applicable_wins_over_all_results() -> None: + """ + Verifies that an inapplicable capability is NOT_APPLICABLE regardless of + results. + """ + # Precondition. + results = [_result(CapabilityCheckStatus.FAILED)] + + # Under test and postcondition. + assert ( + aggregate_capability_verdict(False, results) == CapabilityVerdict.NOT_APPLICABLE + ), "Inapplicable capabilities must ignore check results." + + +def test_required_failure_fails_the_capability() -> None: + """Verifies that any required failure yields FAILED even among passes.""" + # Precondition. + results = [ + _result(CapabilityCheckStatus.PASSED), + _result(CapabilityCheckStatus.FAILED, required=True), + _result(CapabilityCheckStatus.FAILED, required=False), + _result(CapabilityCheckStatus.INDETERMINATE), + ] + + # Under test and postcondition. + assert aggregate_capability_verdict(True, results) == CapabilityVerdict.FAILED + + +def test_only_optional_failures_pass_with_warnings() -> None: + """Verifies that non-required failures downgrade to PASSED_WITH_WARNINGS.""" + # Precondition. + results = [ + _result(CapabilityCheckStatus.PASSED), + _result(CapabilityCheckStatus.FAILED, required=False), + ] + + # Under test and postcondition. + assert ( + aggregate_capability_verdict(True, results) + == CapabilityVerdict.PASSED_WITH_WARNINGS + ) + + +def test_required_indeterminate_outranks_optional_failure() -> None: + """ + Verifies that an unverified required check blocks any PASSED claim, even + when a non-required check produced a definite warning. + """ + # Precondition. + results = [ + _result(CapabilityCheckStatus.FAILED, required=False), + _result(CapabilityCheckStatus.INDETERMINATE, required=True), + ] + + # Under test and postcondition. + assert ( + aggregate_capability_verdict(True, results) == CapabilityVerdict.INDETERMINATE + ) + + +def test_optional_failure_outranks_optional_indeterminate() -> None: + """ + Verifies that with the required core verified, a definite warning outranks + a transient non-required indeterminate. + """ + # Precondition. + results = [ + _result(CapabilityCheckStatus.PASSED, required=True), + _result(CapabilityCheckStatus.FAILED, required=False), + _result(CapabilityCheckStatus.INDETERMINATE, required=False), + ] + + # Under test and postcondition. + assert ( + aggregate_capability_verdict(True, results) + == CapabilityVerdict.PASSED_WITH_WARNINGS + ) + + +def test_indeterminate_without_failures() -> None: + """ + Verifies that a transient error yields INDETERMINATE when nothing failed. + """ + # Precondition. + results = [ + _result(CapabilityCheckStatus.PASSED), + _result(CapabilityCheckStatus.INDETERMINATE), + ] + + # Under test and postcondition. + assert ( + aggregate_capability_verdict(True, results) == CapabilityVerdict.INDETERMINATE + ) + + +def test_all_skipped_yields_skipped() -> None: + """Verifies that a run with only skipped checks is SKIPPED, not PASSED.""" + # Precondition. + results = [ + _result(CapabilityCheckStatus.SKIPPED), + _result(CapabilityCheckStatus.SKIPPED), + ] + + # Under test and postcondition. + assert aggregate_capability_verdict(True, results) == CapabilityVerdict.SKIPPED + + +def test_empty_results_yield_skipped() -> None: + """Verifies that an applicable capability with zero checks is SKIPPED.""" + # Under test and postcondition. + assert aggregate_capability_verdict(True, []) == CapabilityVerdict.SKIPPED + + +def test_skipped_mixed_with_passed_yields_passed() -> None: + """Verifies that partial skips do not mask passing required checks.""" + # Precondition. + results = [ + _result(CapabilityCheckStatus.PASSED), + _result(CapabilityCheckStatus.SKIPPED), + ] + + # Under test and postcondition. + assert aggregate_capability_verdict(True, results) == CapabilityVerdict.PASSED + + +def test_all_passed_yields_passed() -> None: + """Verifies the all-green case.""" + # Precondition. + results = [_result(CapabilityCheckStatus.PASSED)] + + # Under test and postcondition. + assert aggregate_capability_verdict(True, results) == CapabilityVerdict.PASSED + + +def test_compute_capability_verdicts_slack_shaped_scenario() -> None: + """Verifies per-capability grouping for a Slack-shaped report.""" + # Precondition. + # Indexing passes, doc perm sync has a required failure, and group sync is + # not applicable (Slack has none by design). + results = [ + _result(CapabilityCheckStatus.PASSED, capability=CredentialCapability.INDEXING), + _result( + CapabilityCheckStatus.FAILED, + capability=CredentialCapability.DOC_PERMISSION_SYNC, + ), + ] + applicable = { + CredentialCapability.INDEXING, + CredentialCapability.DOC_PERMISSION_SYNC, + } + + # Under test. + verdicts = compute_capability_verdicts(applicable, results) + + # Postcondition. + assert verdicts == { + CredentialCapability.INDEXING: CapabilityVerdict.PASSED, + CredentialCapability.DOC_PERMISSION_SYNC: CapabilityVerdict.FAILED, + CredentialCapability.EXTERNAL_GROUP_SYNC: CapabilityVerdict.NOT_APPLICABLE, + }