From 5a131f4547fc280421e4eb1e6b299eefd97c7818 Mon Sep 17 00:00:00 2001 From: pablonyx Date: Fri, 21 Feb 2025 15:56:11 -0800 Subject: [PATCH] Fix integration tests (#4059) --- .github/workflows/pr-integration-tests.yml | 5 +++-- .../celery/tasks/external_group_syncing/tasks.py | 1 + backend/onyx/background/indexing/run_indexing.py | 5 +++-- backend/onyx/configs/app_configs.py | 4 +++- backend/onyx/connectors/factory.py | 9 ++++++++- backend/onyx/connectors/google_drive/connector.py | 9 ++++++++- backend/onyx/db/connector_credential_pair.py | 5 +++++ backend/tests/integration/common_utils/constants.py | 2 +- .../tests/integration/common_utils/managers/connector.py | 6 ++++-- backend/tests/integration/common_utils/managers/user.py | 2 -- deployment/docker_compose/docker-compose.dev.yml | 1 + 11 files changed, 37 insertions(+), 12 deletions(-) diff --git a/.github/workflows/pr-integration-tests.yml b/.github/workflows/pr-integration-tests.yml index 5573c51d9f2..593a9adfd24 100644 --- a/.github/workflows/pr-integration-tests.yml +++ b/.github/workflows/pr-integration-tests.yml @@ -145,7 +145,7 @@ jobs: run: | cd deployment/docker_compose docker compose -f docker-compose.multitenant-dev.yml -p onyx-stack down -v - + # NOTE: Use pre-ping/null pool to reduce flakiness due to dropped connections - name: Start Docker containers run: | @@ -157,6 +157,7 @@ jobs: REQUIRE_EMAIL_VERIFICATION=false \ DISABLE_TELEMETRY=true \ IMAGE_TAG=test \ + INTEGRATION_TESTS_MODE=true \ docker compose -f docker-compose.dev.yml -p onyx-stack up -d id: start_docker @@ -199,7 +200,7 @@ jobs: cd backend/tests/integration/mock_services docker compose -f docker-compose.mock-it-services.yml \ -p mock-it-services-stack up -d - + # NOTE: Use pre-ping/null to reduce flakiness due to dropped connections - name: Run Standard Integration Tests run: | diff --git a/backend/onyx/background/celery/tasks/external_group_syncing/tasks.py b/backend/onyx/background/celery/tasks/external_group_syncing/tasks.py index 50bdb5f32e4..2db26cfb6d0 100644 --- a/backend/onyx/background/celery/tasks/external_group_syncing/tasks.py +++ b/backend/onyx/background/celery/tasks/external_group_syncing/tasks.py @@ -361,6 +361,7 @@ def connector_external_group_sync_generator_task( cc_pair = get_connector_credential_pair_from_id( db_session=db_session, cc_pair_id=cc_pair_id, + eager_load_credential=True, ) if cc_pair is None: raise ValueError( diff --git a/backend/onyx/background/indexing/run_indexing.py b/backend/onyx/background/indexing/run_indexing.py index 41823754a94..9729f61dd32 100644 --- a/backend/onyx/background/indexing/run_indexing.py +++ b/backend/onyx/background/indexing/run_indexing.py @@ -15,6 +15,7 @@ from onyx.configs.app_configs import INDEX_BATCH_SIZE from onyx.configs.app_configs import INDEXING_SIZE_WARNING_THRESHOLD from onyx.configs.app_configs import INDEXING_TRACER_INTERVAL +from onyx.configs.app_configs import INTEGRATION_TESTS_MODE from onyx.configs.app_configs import LEAVE_CONNECTOR_ACTIVE_ON_INITIALIZATION_FAILURE from onyx.configs.app_configs import POLL_CONNECTOR_OFFSET from onyx.configs.constants import DocumentSource @@ -89,8 +90,8 @@ def _get_connector_runner( ) # validate the connector settings - - runnable_connector.validate_connector_settings() + if not INTEGRATION_TESTS_MODE: + runnable_connector.validate_connector_settings() except Exception as e: logger.exception(f"Unable to instantiate connector due to {e}") diff --git a/backend/onyx/configs/app_configs.py b/backend/onyx/configs/app_configs.py index 5e87734d222..ec0e4e2d561 100644 --- a/backend/onyx/configs/app_configs.py +++ b/backend/onyx/configs/app_configs.py @@ -158,7 +158,7 @@ POSTGRES_PASSWORD = urllib.parse.quote_plus( os.environ.get("POSTGRES_PASSWORD") or "password" ) -POSTGRES_HOST = os.environ.get("POSTGRES_HOST") or "localhost" +POSTGRES_HOST = os.environ.get("POSTGRES_HOST") or "127.0.0.1" POSTGRES_PORT = os.environ.get("POSTGRES_PORT") or "5432" POSTGRES_DB = os.environ.get("POSTGRES_DB") or "postgres" AWS_REGION_NAME = os.environ.get("AWS_REGION_NAME") or "us-east-2" @@ -626,6 +626,8 @@ DEV_MODE = os.environ.get("DEV_MODE", "").lower() == "true" +INTEGRATION_TESTS_MODE = os.environ.get("INTEGRATION_TESTS_MODE", "").lower() == "true" + MOCK_CONNECTOR_FILE_PATH = os.environ.get("MOCK_CONNECTOR_FILE_PATH") TEST_ENV = os.environ.get("TEST_ENV", "").lower() == "true" diff --git a/backend/onyx/connectors/factory.py b/backend/onyx/connectors/factory.py index b4f497f65dd..c67b08ff86c 100644 --- a/backend/onyx/connectors/factory.py +++ b/backend/onyx/connectors/factory.py @@ -3,6 +3,7 @@ from sqlalchemy.orm import Session +from onyx.configs.app_configs import INTEGRATION_TESTS_MODE from onyx.configs.constants import DocumentSource from onyx.configs.constants import DocumentSourceRequiringTenantContext from onyx.connectors.airtable.airtable_connector import AirtableConnector @@ -187,6 +188,9 @@ def validate_ccpair_for_user( user: User | None, tenant_id: str | None, ) -> None: + if INTEGRATION_TESTS_MODE: + return + # Validate the connector settings connector = fetch_connector_by_id(connector_id, db_session) credential = fetch_credential_by_id_for_user( @@ -199,7 +203,10 @@ def validate_ccpair_for_user( if not connector: raise ValueError("Connector not found") - if connector.source == DocumentSource.INGESTION_API: + if ( + connector.source == DocumentSource.INGESTION_API + or connector.source == DocumentSource.MOCK_CONNECTOR + ): return if not credential: diff --git a/backend/onyx/connectors/google_drive/connector.py b/backend/onyx/connectors/google_drive/connector.py index 1287a896076..46da04ff3e3 100644 --- a/backend/onyx/connectors/google_drive/connector.py +++ b/backend/onyx/connectors/google_drive/connector.py @@ -220,7 +220,14 @@ def creds(self) -> OAuthCredentials | ServiceAccountCredentials: return self._creds def load_credentials(self, credentials: dict[str, Any]) -> dict[str, str] | None: - self._primary_admin_email = credentials[DB_CREDENTIALS_PRIMARY_ADMIN_KEY] + try: + self._primary_admin_email = credentials[DB_CREDENTIALS_PRIMARY_ADMIN_KEY] + except KeyError: + raise ValueError( + "Primary admin email missing, " + "should not call this property " + "before calling load_credentials" + ) self._creds, new_creds_dict = get_google_creds( credentials=credentials, diff --git a/backend/onyx/db/connector_credential_pair.py b/backend/onyx/db/connector_credential_pair.py index 712e81894a0..7c73faaa7b9 100644 --- a/backend/onyx/db/connector_credential_pair.py +++ b/backend/onyx/db/connector_credential_pair.py @@ -194,9 +194,14 @@ def get_connector_credential_pair_from_id_for_user( def get_connector_credential_pair_from_id( db_session: Session, cc_pair_id: int, + eager_load_credential: bool = False, ) -> ConnectorCredentialPair | None: stmt = select(ConnectorCredentialPair).distinct() stmt = stmt.where(ConnectorCredentialPair.id == cc_pair_id) + + if eager_load_credential: + stmt = stmt.options(joinedload(ConnectorCredentialPair.credential)) + result = db_session.execute(stmt) return result.scalar_one_or_none() diff --git a/backend/tests/integration/common_utils/constants.py b/backend/tests/integration/common_utils/constants.py index c6731e7397a..5f0247852ab 100644 --- a/backend/tests/integration/common_utils/constants.py +++ b/backend/tests/integration/common_utils/constants.py @@ -3,7 +3,7 @@ ADMIN_USER_NAME = "admin_user" API_SERVER_PROTOCOL = os.getenv("API_SERVER_PROTOCOL") or "http" -API_SERVER_HOST = os.getenv("API_SERVER_HOST") or "localhost" +API_SERVER_HOST = os.getenv("API_SERVER_HOST") or "127.0.0.1" API_SERVER_PORT = os.getenv("API_SERVER_PORT") or "8080" API_SERVER_URL = f"{API_SERVER_PROTOCOL}://{API_SERVER_HOST}:{API_SERVER_PORT}" MAX_DELAY = 45 diff --git a/backend/tests/integration/common_utils/managers/connector.py b/backend/tests/integration/common_utils/managers/connector.py index 04dd37c2adf..410182a9219 100644 --- a/backend/tests/integration/common_utils/managers/connector.py +++ b/backend/tests/integration/common_utils/managers/connector.py @@ -30,8 +30,10 @@ def create( name=name, source=source, input_type=input_type, - connector_specific_config=connector_specific_config - or {"file_locations": []}, + connector_specific_config=( + connector_specific_config + or ({"file_locations": []} if source == DocumentSource.FILE else {}) + ), access_type=access_type, groups=groups or [], ) diff --git a/backend/tests/integration/common_utils/managers/user.py b/backend/tests/integration/common_utils/managers/user.py index 0045e799522..950336f11d6 100644 --- a/backend/tests/integration/common_utils/managers/user.py +++ b/backend/tests/integration/common_utils/managers/user.py @@ -88,8 +88,6 @@ def login_as_user(test_user: DATestUser) -> DATestUser: if not session_cookie: raise Exception("Failed to login") - print(f"Logged in as {test_user.email}") - # Set cookies in the headers test_user.headers["Cookie"] = f"fastapiusersauth={session_cookie}; " test_user.cookies = {"fastapiusersauth": session_cookie} diff --git a/deployment/docker_compose/docker-compose.dev.yml b/deployment/docker_compose/docker-compose.dev.yml index 14860827c5d..525246e54ed 100644 --- a/deployment/docker_compose/docker-compose.dev.yml +++ b/deployment/docker_compose/docker-compose.dev.yml @@ -36,6 +36,7 @@ services: - OPENID_CONFIG_URL=${OPENID_CONFIG_URL:-} - TRACK_EXTERNAL_IDP_EXPIRY=${TRACK_EXTERNAL_IDP_EXPIRY:-} - CORS_ALLOWED_ORIGIN=${CORS_ALLOWED_ORIGIN:-} + - INTEGRATION_TESTS_MODE=${INTEGRATION_TESTS_MODE:-} # Gen AI Settings - GEN_AI_MAX_TOKENS=${GEN_AI_MAX_TOKENS:-} - QA_TIMEOUT=${QA_TIMEOUT:-}