Abstract os.environ so real env doesn't interfere with tests

src/buildkite_test_collector/collector/api.py

@@ -1,45 +1,51 @@
 """Buildkite Test Analytics API"""
 
-from typing import Any, Generator, Optional
-from os import environ
+from typing import Any, Generator, Optional, Mapping
 import traceback
 from requests import post, Response
 from requests.exceptions import InvalidHeader, HTTPError
 from .payload import Payload
 from ..pytest_plugin.logger import logger
 
 
-def submit(payload: Payload, batch_size=100) -> Generator[Optional[Response], Any, Any]:
-    """Submit a payload to the API"""
-    token = environ.get("BUILDKITE_ANALYTICS_TOKEN")
-    api_url = environ.get("BUILDKITE_ANALYTICS_API_URL", "https://analytics-api.buildkite.com/v1")
-    response = None
+class API:
+    """Buildkite Test Analytics API client"""
 
-    if not token:
-        logger.warning("No `BUILDKITE_ANALYTICS_TOKEN` environment variable present")
-        yield None
+    def __init__(self, env: Mapping[str, Optional[str]]):
+        """Initialize the API client with environment variables"""
+        self.env = env
 
-    else:
-        for payload_slice in payload.into_batches(batch_size):
-            try:
-                response = post(api_url + "/uploads",
-                                json=payload_slice.as_json(),
-                                headers={
-                                    "Content-Type": "application/json",
-                                    "Authorization": f"Token token=\"{token}\""
-                                },
-                                timeout=60)
-                response.raise_for_status()
-                yield response
-            except InvalidHeader as error:
-                logger.warning("Invalid `BUILDKITE_ANALYTICS_TOKEN` environment variable")
-                logger.warning(error)
-                yield None
-            except HTTPError as err:
-                logger.warning("Failed to uploads test results to buildkite")
-                logger.warning(err)
-                yield None
-            except Exception:  # pylint: disable=broad-except
-                error_message = traceback.format_exc()
-                logger.warning(error_message)
-                yield None
+    def submit(self, payload: Payload, batch_size=100) -> Generator[Optional[Response], Any, Any]:
+        """Submit a payload to the API"""
+        token = self.env.get("BUILDKITE_ANALYTICS_TOKEN")
+        api_url = self.env.get("BUILDKITE_ANALYTICS_API_URL") or "https://analytics-api.buildkite.com/v1"
+        response = None
+
+        if not token:
+            logger.warning("No `BUILDKITE_ANALYTICS_TOKEN` environment variable present")
+            yield None
+
+        else:
+            for payload_slice in payload.into_batches(batch_size):
+                try:
+                    response = post(api_url + "/uploads",
+                                    json=payload_slice.as_json(),
+                                    headers={
+                                        "Content-Type": "application/json",
+                                        "Authorization": f"Token token=\"{token}\""
+                                    },
+                                    timeout=60)
+                    response.raise_for_status()
+                    yield response
+                except InvalidHeader as error:
+                    logger.warning("Invalid `BUILDKITE_ANALYTICS_TOKEN` environment variable")
+                    logger.warning(error)
+                    yield None
+                except HTTPError as err:
+                    logger.warning("Failed to uploads test results to buildkite")
+                    logger.warning(err)
+                    yield None
+                except Exception:  # pylint: disable=broad-except
+                    error_message = traceback.format_exc()
+                    logger.warning(error_message)
+                    yield None

src/buildkite_test_collector/collector/run_env.py

@@ -1,93 +1,119 @@
 """Test Engine run_env"""
 
 from dataclasses import dataclass
-from typing import Dict, Optional
+from typing import Dict, Optional, Mapping
 from uuid import uuid4
-import os
-from .constants import COLLECTOR_NAME, VERSION # pylint: disable=W0611
-
-# pylint: disable=C0103 disable=R0902
-
-
-def __get_env(name: str) -> Optional[str]:
-    var = os.environ.get(name)
-    if (var is None or var == ''):
-        return None
-
-    return var
-
-
-def __buildkite_env() -> Optional['RunEnv']:
-    build_id = __get_env("BUILDKITE_BUILD_ID")
-
-    if build_id is None:
-        return None
-
-    return RunEnv(
-        ci="buildkite",
-        key=build_id,
-        url=__get_env("BUILDKITE_BUILD_URL"),
-        branch=__get_env("BUILDKITE_BRANCH"),
-        commit_sha=__get_env("BUILDKITE_COMMIT"),
-        number=__get_env("BUILDKITE_BUILD_NUMBER"),
-        job_id=__get_env("BUILDKITE_JOB_ID"),
-        message=__get_env("BUILDKITE_MESSAGE")
-    )
-
 
-def __github_actions_env() -> Optional['RunEnv']:
-    action = __get_env("GITHUB_ACTION")
-    run_number = __get_env("GITHUB_RUN_NUMBER")
-    run_attempt = __get_env("GITHUB_RUN_ATTEMPT")
-
-    if (action is None or run_number is None or run_attempt is None):
-        return None
-
-    repo = __get_env("GITHUB_REPOSITORY")
-    run_id = __get_env("GITHUB_RUN_ID")
-
-    return RunEnv(
-        ci="github_actions",
-        key=f"{action}-{run_number}-{run_attempt}",
-        url=f"https://github.com/{repo}/actions/runs/{run_id}",
-        branch=__get_env("GITHUB_REF"),
-        commit_sha=__get_env("GITHUB_SHA"),
-        number=run_number,
-        job_id=None,
-        message=__get_env("TEST_ANALYTICS_COMMIT_MESSAGE"),
-    )
-
-
-def __circle_ci_env() -> Optional['RunEnv']:
-    build_num = __get_env("CIRCLE_BUILD_NUM")
-    workflow_id = __get_env("CIRCLE_WORKFLOW_ID")
+from .constants import COLLECTOR_NAME, VERSION # pylint: disable=W0611
 
-    if (build_num is None or workflow_id is None):
-        return None
+# pylint: disable=R0902
 
-    return RunEnv(
-        ci="circleci",
-        key=f"{workflow_id}-{build_num}",
-        url=__get_env("CIRCLE_BUILD_URL"),
-        branch=__get_env("CIRCLE_BRANCH"),
-        commit_sha=__get_env("CIRCLE_SHA1"),
-        number=build_num,
-        job_id=None,
-        message=__get_env("TEST_ANALYTICS_COMMIT_MESSAGE"),
-    )
+class RunEnvBuilder:
+    """Builder class for RunEnv that allows injection of environment variables
 
+    Example usage:
+        # Normal usage
+        builder = RunEnvBuilder(os.environ)
+        env = builder.build()
 
-def __generic_env() -> 'RunEnv':
-    return RunEnv(
-        ci="generic",
-        key=str(uuid4()),
-        url=None,
-        branch=None,
-        commit_sha=None,
-        number=None,
-        job_id=None,
-        message=None
-    )
+        # Testing with fake environment
+        fake_env = {
+            "BUILDKITE_BUILD_ID": "test-build-123",
+            "BUILDKITE_BUILD_URL": "https://buildkite.com/test",
+            "BUILDKITE_BRANCH": "main",
+            "BUILDKITE_COMMIT": "abc123",
+        }
+        builder = RunEnvBuilder(fake_env)
+        env = builder.build()
+        assert env.ci == "buildkite"
+        assert env.key == "test-build-123"
+    """
+
+    def __init__(self, env: Mapping[str, Optional[str]]):
+        self.env = env
+
+    def build(self) -> 'RunEnv':
+        """Build a RunEnv by detecting the CI system"""
+        return \
+            self._buildkite_env() or \
+            self._github_actions_env() or \
+            self._circle_ci_env() or \
+            self._generic_env()
+
+    def _get_env(self, name: str) -> Optional[str]:
+        var = self.env.get(name)
+        if (var is None or var == ''):
+            return None
+        return var
+
+    def _buildkite_env(self) -> Optional['RunEnv']:
+        build_id = self._get_env("BUILDKITE_BUILD_ID")
+
+        if build_id is None:
+            return None
+
+        return RunEnv(
+            ci="buildkite",
+            key=build_id,
+            url=self._get_env("BUILDKITE_BUILD_URL"),
+            branch=self._get_env("BUILDKITE_BRANCH"),
+            commit_sha=self._get_env("BUILDKITE_COMMIT"),
+            number=self._get_env("BUILDKITE_BUILD_NUMBER"),
+            job_id=self._get_env("BUILDKITE_JOB_ID"),
+            message=self._get_env("BUILDKITE_MESSAGE")
+        )
+
+    def _github_actions_env(self) -> Optional['RunEnv']:
+        action = self._get_env("GITHUB_ACTION")
+        run_number = self._get_env("GITHUB_RUN_NUMBER")
+        run_attempt = self._get_env("GITHUB_RUN_ATTEMPT")
+
+        if (action is None or run_number is None or run_attempt is None):
+            return None
+
+        repo = self._get_env("GITHUB_REPOSITORY")
+        run_id = self._get_env("GITHUB_RUN_ID")
+
+        return RunEnv(
+            ci="github_actions",
+            key=f"{action}-{run_number}-{run_attempt}",
+            url=f"https://github.com/{repo}/actions/runs/{run_id}",
+            branch=self._get_env("GITHUB_REF"),
+            commit_sha=self._get_env("GITHUB_SHA"),
+            number=run_number,
+            job_id=None,
+            message=self._get_env("TEST_ANALYTICS_COMMIT_MESSAGE"),
+        )
+
+    def _circle_ci_env(self) -> Optional['RunEnv']:
+        build_num = self._get_env("CIRCLE_BUILD_NUM")
+        workflow_id = self._get_env("CIRCLE_WORKFLOW_ID")
+
+        if (build_num is None or workflow_id is None):
+            return None
+
+        return RunEnv(
+            ci="circleci",
+            key=f"{workflow_id}-{build_num}",
+            url=self._get_env("CIRCLE_BUILD_URL"),
+            branch=self._get_env("CIRCLE_BRANCH"),
+            commit_sha=self._get_env("CIRCLE_SHA1"),
+            number=build_num,
+            job_id=None,
+            message=self._get_env("TEST_ANALYTICS_COMMIT_MESSAGE"),
+        )
+
+    def _generic_env(self) -> 'RunEnv':
+        return RunEnv(
+            ci="generic",
+            key=str(uuid4()),
+            url=None,
+            branch=None,
+            commit_sha=None,
+            number=None,
+            job_id=None,
+            message=None
+        )
 
 
 @dataclass(frozen=True)
@@ -118,11 +144,3 @@ def as_json(self) -> Dict[str, str]:
         }
 
         return {k: v for k, v in attrs.items() if v is not None}
-
-
-def detect_env() -> RunEnv:
-    """Attempt to detect the CI system we're running in"""
-    return __buildkite_env() or \
-        __github_actions_env() or \
-        __circle_ci_env() or \
-        __generic_env()

src/buildkite_test_collector/pytest_plugin/__init__.py

@@ -1,11 +1,11 @@
-# pylint: disable=line-too-long
 """Buildkite test collector for Pytest."""
 
+import os
 import pytest
 
 from ..collector.payload import Payload
-from ..collector.run_env import detect_env
-from ..collector.api import submit
+from ..collector.run_env import RunEnvBuilder
+from ..collector.api import API
 from .span_collector import SpanCollector
 from .buildkite_plugin import BuildkitePlugin
 
@@ -22,9 +22,12 @@ def spans(request):
 @pytest.hookimpl(trylast=True)
 def pytest_configure(config):
     """pytest_configure hook callback"""
-    env = detect_env()
+    env = RunEnvBuilder(os.environ).build()
 
-    config.addinivalue_line("markers", "execution_tag(key, value): add tag to test execution for Buildkite Test Collector. Both key and value must be a string.")
+    config.addinivalue_line("markers",
+        "execution_tag(key, value): "
+        "add tag to test execution for Buildkite Test Collector. "
+        "Both key and value must be a string.")
 
     plugin = BuildkitePlugin(Payload.init(env))
     setattr(config, '_buildkite', plugin)
@@ -37,6 +40,7 @@ def pytest_unconfigure(config):
     plugin = getattr(config, '_buildkite', None)
 
     if plugin:
+        api = API(os.environ)
         xdist_enabled = (
             config.pluginmanager.getplugin("xdist") is not None
             and config.getoption("numprocesses") is not None
@@ -47,11 +51,12 @@ def pytest_unconfigure(config):
 
         # When xdist is not installed, or when it's installed and not enabled
         if not xdist_enabled:
-            list(submit(plugin.payload))
+            list(api.submit(plugin.payload))
 
-        # When xdist is activated, we want to submit from worker thread only, because they have access to tag data
+        # When xdist is activated, we want to submit from worker thread only, because they have
+        # access to tag data
         if xdist_enabled and is_xdist_worker:
-            list(submit(plugin.payload))
+            list(api.submit(plugin.payload))
 
         # We only want a single thread to write to the json file.
         # When xdist is enabled, that will be the controller thread.