refactor(smoke test): centralise env variables (#15100)

Author: anshbansal

smoke-test/conftest.py

@@ -7,6 +7,7 @@
 from datahub.ingestion.graph.client import DatahubClientConfig, DataHubGraph, get_default_graph
 
 from tests.test_result_msg import send_message
+from tests.utilities import env_vars
 from tests.utils import (
     TestSessionWrapper,
     get_frontend_session,
@@ -149,11 +150,9 @@ def bin_pack_tasks(tasks, n_buckets):
     return buckets
 
 def get_batch_start_end(num_tests: int) -> Tuple[int, int]:
-    batch_count_env = os.getenv("BATCH_COUNT", 1)
-    batch_count = int(batch_count_env)
+    batch_count = env_vars.get_batch_count()
 
-    batch_number_env = os.getenv("BATCH_NUMBER", 0)
-    batch_number = int(batch_number_env)
+    batch_number = env_vars.get_batch_number()
 
     if batch_count == 0 or batch_count > num_tests:
         raise ValueError(
@@ -182,7 +181,7 @@ def get_batch_start_end(num_tests: int) -> Tuple[int, int]:
 def pytest_collection_modifyitems(
     session: pytest.Session, config: pytest.Config, items: List[Item]
 ) -> None:
-    if os.getenv("TEST_STRATEGY") == "cypress":
+    if env_vars.get_test_strategy() == "cypress":
         return  # We launch cypress via pytests, but needs a different batching mechanism at cypress level.
 
     # If BATCH_COUNT and BATCH_ENV vars are set, splits the pytests to batches and runs filters only the BATCH_NUMBER

smoke-test/tests/actions/doc_propagation/test_propagation.py

@@ -21,6 +21,7 @@
 from datahub.ingestion.graph.client import DatahubClientConfig, DataHubGraph
 from datahub.ingestion.sink.file import FileSink, FileSinkConfig
 from datahub.utilities.urns.urn import Urn
+from tests.utilities import env_vars
 from tests.utils import (
     delete_urns_from_file,
     get_gms_url,
@@ -30,7 +31,7 @@
 
 logger = logging.getLogger(__name__)
 
-DELETE_AFTER_TEST = os.getenv("DELETE_AFTER_TEST", "false").lower() == "true"
+DELETE_AFTER_TEST = env_vars.get_delete_after_test()
 
 
 class FileEmitter:

smoke-test/tests/consistency_utils.py

@@ -1,13 +1,14 @@
 import logging
-import os
 import subprocess
 import time
 
-USE_STATIC_SLEEP: bool = bool(os.getenv("USE_STATIC_SLEEP", False))
-ELASTICSEARCH_REFRESH_INTERVAL_SECONDS: int = int(
-    os.getenv("ELASTICSEARCH_REFRESH_INTERVAL_SECONDS", 3)
+from tests.utilities import env_vars
+
+USE_STATIC_SLEEP: bool = env_vars.get_use_static_sleep()
+ELASTICSEARCH_REFRESH_INTERVAL_SECONDS: int = (
+    env_vars.get_elasticsearch_refresh_interval_seconds()
 )
-KAFKA_BOOTSTRAP_SERVER: str = str(os.getenv("KAFKA_BOOTSTRAP_SERVER", "broker:29092"))
+KAFKA_BOOTSTRAP_SERVER: str = env_vars.get_kafka_bootstrap_server()
 
 logger = logging.getLogger(__name__)
 
@@ -35,7 +36,7 @@ def wait_for_writes_to_sync(
         time.sleep(ELASTICSEARCH_REFRESH_INTERVAL_SECONDS)
         return
     KAFKA_BROKER_CONTAINER: str = str(
-        os.getenv("KAFKA_BROKER_CONTAINER", infer_kafka_broker_container())
+        env_vars.get_kafka_broker_container() or infer_kafka_broker_container()
     )
     start_time = time.time()
     # get offsets

smoke-test/tests/cypress/integration_test.py

@@ -13,6 +13,7 @@
     get_time_lineage_urns,
     ingest_time_lineage,
 )
+from tests.utilities import env_vars
 from tests.utils import (
     create_datahub_step_state_aspects,
     delete_urns,
@@ -234,15 +235,15 @@ def _get_cypress_tests_batch():
         else:
             tests_with_weights.append(test)
 
-    test_batches = bin_pack_tasks(tests_with_weights, int(os.getenv("BATCH_COUNT", 1)))
-    return test_batches[int(os.getenv("BATCH_NUMBER", 0))]
+    test_batches = bin_pack_tasks(tests_with_weights, env_vars.get_batch_count())
+    return test_batches[env_vars.get_batch_number()]
 
 
 def test_run_cypress(auth_session):
     # Run with --record option only if CYPRESS_RECORD_KEY is non-empty
-    record_key = os.getenv("CYPRESS_RECORD_KEY")
+    record_key = env_vars.get_cypress_record_key()
     tag_arg = ""
-    test_strategy = os.getenv("TEST_STRATEGY", None)
+    test_strategy = env_vars.get_test_strategy()
     if record_key:
         record_arg = " --record "
         batch_number = os.getenv("BATCH_NUMBER")

smoke-test/tests/openapi/v1/test_tracking.py

@@ -15,7 +15,6 @@
 """
 
 import json
-import os
 import time
 from datetime import datetime, timezone
 
@@ -24,6 +23,7 @@
 from confluent_kafka import Consumer, KafkaError
 from dotenv import load_dotenv
 
+from tests.utilities import env_vars
 from tests.utils import get_kafka_broker_url
 
 # Load environment variables from .env file if it exists
@@ -37,7 +37,7 @@ def test_tracking_api_mixpanel(auth_session, graph_client):
     """Test that we can post events to the tracking endpoint and verify they are sent to Mixpanel."""
 
     # Check if Mixpanel API secret is available in environment variables
-    api_secret = os.environ.get("MIXPANEL_API_SECRET")
+    api_secret = env_vars.get_mixpanel_api_secret()
     if not api_secret:
         pytest.skip("MIXPANEL_API_SECRET environment variable not set, skipping test")
 
@@ -80,9 +80,7 @@ def test_tracking_api_mixpanel(auth_session, graph_client):
 
     # Query Mixpanel's JQL API to retrieve our test event
     # Note: This requires a service account with access to JQL
-    project_id = os.environ.get(
-        "MIXPANEL_PROJECT_ID", "3653440"
-    )  # Allow project ID to be configurable too
+    project_id = env_vars.get_mixpanel_project_id()
 
     # log the unique_id
     print(f"\nLooking for test event with customField: {unique_id}")
@@ -357,8 +355,8 @@ def test_tracking_api_elasticsearch(auth_session):
 
     # Query Elasticsearch to retrieve our test event
     # This requires the Elasticsearch URL and credentials
-    es_url = os.environ.get("ELASTICSEARCH_URL", "http://localhost:9200")
-    es_index = os.environ.get("ELASTICSEARCH_INDEX", "datahub_usage_event")
+    es_url = env_vars.get_elasticsearch_url()
+    es_index = env_vars.get_elasticsearch_index()
 
     # Create a query to find our test event by the unique browserId and customField
     es_query = {

smoke-test/tests/read_only/test_services_up.py

@@ -1,10 +1,11 @@
-import os
 import re
 
 import pytest
 
+from tests.utilities import env_vars
+
 # Kept separate so that it does not cause failures in PRs
-DATAHUB_VERSION = os.getenv("TEST_DATAHUB_VERSION")
+DATAHUB_VERSION = env_vars.get_test_datahub_version()
 
 
 def looks_like_a_short_sha(sha: str) -> bool:

smoke-test/tests/test_result_msg.py

@@ -1,7 +1,7 @@
-import os
-
 from slack_sdk import WebClient
 
+from tests.utilities import env_vars
+
 datahub_stats = {}
 
 
@@ -10,10 +10,10 @@ def add_datahub_stats(stat_name, stat_val):
 
 
 def send_to_slack(passed: str):
-    slack_api_token = os.getenv("SLACK_API_TOKEN")
-    slack_channel = os.getenv("SLACK_CHANNEL")
-    slack_thread_ts = os.getenv("SLACK_THREAD_TS")
-    test_identifier = os.getenv("TEST_IDENTIFIER", "LOCAL_TEST")
+    slack_api_token = env_vars.get_slack_api_token()
+    slack_channel = env_vars.get_slack_channel()
+    slack_thread_ts = env_vars.get_slack_thread_ts()
+    test_identifier = env_vars.get_test_identifier()
     if slack_api_token is None or slack_channel is None:
         return
     client = WebClient(token=slack_api_token)