use env var

datahub-project · Jul 28, 2022 · d43af50 · d43af50
1 parent 25a9a03
commit d43af50
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 19 deletions.
diff --git a/metadata-ingestion/src/datahub/emitter/rest_emitter.py b/metadata-ingestion/src/datahub/emitter/rest_emitter.py
@@ -1,6 +1,7 @@
 import datetime
 import json
 import logging
+import os
 from json.decoder import JSONDecodeError
 from typing import Any, Dict, List, Optional, Tuple, Union
 
@@ -34,7 +35,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
@@ -1,26 +1,31 @@
 import logging
 import os
 import time
-import unittest.mock
 
 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.
     import pandas  # noqa: F401
 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):
@@ -32,14 +37,6 @@ def fake_time():
     yield
 
 
-@pytest.fixture(autouse=True, scope="session")
-def reduce_gms_retries():
-    with unittest.mock.patch(
-        "datahub.emitter.rest_emitter.DataHubRestEmitter._retry_max_times", 1
-    ):
-        yield
-
-
 def pytest_addoption(parser):
     parser.addoption(
         "--update-golden-files",