diff --git a/metadata-ingestion/src/datahub/emitter/rest_emitter.py b/metadata-ingestion/src/datahub/emitter/rest_emitter.py index 2360c1c66f4ca7..21abdd80d9c77b 100644 --- a/metadata-ingestion/src/datahub/emitter/rest_emitter.py +++ b/metadata-ingestion/src/datahub/emitter/rest_emitter.py @@ -2,6 +2,7 @@ import functools import json import logging +import os from json.decoder import JSONDecodeError from typing import Any, Dict, List, Optional, Tuple, Union @@ -35,7 +36,9 @@ class DataHubRestEmitter: 504, ] DEFAULT_RETRY_METHODS = ["HEAD", "GET", "POST", "PUT", "DELETE", "OPTIONS", "TRACE"] - DEFAULT_RETRY_MAX_TIMES = 3 + DEFAULT_RETRY_MAX_TIMES = int( + os.getenv("DATAHUB_REST_EMITTER_DEFAULT_RETRY_MAX_TIMES", "3") + ) _gms_server: str _token: Optional[str] diff --git a/metadata-ingestion/tests/conftest.py b/metadata-ingestion/tests/conftest.py index 800232c921c640..0390547800b765 100644 --- a/metadata-ingestion/tests/conftest.py +++ b/metadata-ingestion/tests/conftest.py @@ -4,8 +4,21 @@ import pytest -from tests.test_helpers.docker_helpers import docker_compose_runner # noqa: F401 -from tests.test_helpers.state_helpers import mock_datahub_graph # noqa: F401 +# Enable debug logging. +logging.getLogger().setLevel(logging.DEBUG) +os.environ["DATAHUB_DEBUG"] = "1" + +# Disable telemetry +os.environ["DATAHUB_TELEMETRY_ENABLED"] = "false" + +# Reduce retries on GMS, because this causes tests to hang while sleeping +# between retries. +os.environ["DATAHUB_REST_EMITTER_DEFAULT_RETRY_MAX_TIMES"] = "1" + +# We need our imports to go below the os.environ updates, since mere act +# of importing some datahub modules will load env variables. +from tests.test_helpers.docker_helpers import docker_compose_runner # noqa: F401,E402 +from tests.test_helpers.state_helpers import mock_datahub_graph # noqa: F401,E402 try: # See https://github.com/spulec/freezegun/issues/98#issuecomment-590553475. @@ -13,13 +26,6 @@ except ImportError: pass -# Enable debug logging. -logging.getLogger().setLevel(logging.DEBUG) -os.putenv("DATAHUB_DEBUG", "1") - -# Disable telemetry -os.putenv("DATAHUB_TELEMETRY_ENABLED", "false") - @pytest.fixture def mock_time(monkeypatch):