Fix telemetry issues in wheel build and unit tests (pybamm-team#4591)

* Reorganize to make sure capture does not run

* Cleanup

* Mock telemetry object

* Fix logic

* Minor fixes for networking

* temporary trigger

* testing env

* Remove requirement

* Try again

* Last attempt

* Update wheel env

* Cleanup

* More cleanup

* Final test for cleanup

* Turn off wheel build

* Make sure opt out is in all tests

* Fix some coverage

* Cover last line

* style: pre-commit fixes

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

.github/workflows/benchmark_on_push.yml

@@ -8,6 +8,9 @@ concurrency:
   group: ${{ github.workflow }}-${{ github.ref }}
   cancel-in-progress: true
 
+env:
+  PYBAMM_DISABLE_TELEMETRY: "true"
+
 jobs:
   benchmarks:
     runs-on: ubuntu-latest

.github/workflows/periodic_benchmarks.yml

@@ -15,6 +15,9 @@ on:
   # workflow manually
   workflow_dispatch:
 
+env:
+  PYBAMM_DISABLE_TELEMETRY: "true"
+
 jobs:
   benchmarks:
     runs-on: ubuntu-latest

.github/workflows/publish_pypi.yml

@@ -18,6 +18,7 @@ on:
 
 # Set options available for all jobs that use cibuildwheel
 env:
+  PYBAMM_DISABLE_TELEMETRY: "true"
   # Increase pip debugging output, equivalent to `pip -vv`
   CIBW_BUILD_VERBOSITY: 2
   # Disable build isolation to allow pre-installing build-time dependencies.
@@ -75,6 +76,7 @@ jobs:
         run: pipx run cibuildwheel --output-dir wheelhouse
         env:
           CIBW_ENVIRONMENT: >
+            PYBAMM_DISABLE_TELEMETRY="true"
             PYBAMM_USE_VCPKG=ON
             VCPKG_ROOT_DIR=C:\vcpkg
             VCPKG_DEFAULT_TRIPLET=x64-windows-static-md
@@ -116,6 +118,8 @@ jobs:
       - name: Build wheels on Linux
         run: pipx run cibuildwheel --output-dir wheelhouse
         env:
+          CIBW_ENVIRONMENT: >
+            PYBAMM_DISABLE_TELEMETRY="true"
           CIBW_ARCHS_LINUX: x86_64
           CIBW_BEFORE_ALL_LINUX: >
             yum -y install openblas-devel lapack-devel &&
@@ -242,6 +246,8 @@ jobs:
             python scripts/install_KLU_Sundials.py
             python -m cibuildwheel --output-dir wheelhouse
         env:
+          CIBW_ENVIRONMENT: >
+            PYBAMM_DISABLE_TELEMETRY="true"
           # 10.13 for Intel (macos-13), 11.0 for Apple Silicon (macos-14 and macos-latest)
           MACOSX_DEPLOYMENT_TARGET: ${{ matrix.os == 'macos-14' && '11.0' || '10.13' }}
           CIBW_ARCHS_MACOS: auto

.github/workflows/run_benchmarks_over_history.yml

@@ -18,6 +18,10 @@ on:
       ncommits:
         description: "Number of commits to benchmark between commit_start and commit_end"
         default: "100"
+
+env:
+  PYBAMM_DISABLE_TELEMETRY: "true"
+
 jobs:
   benchmarks:
     runs-on: ubuntu-latest

.github/workflows/run_periodic_tests.yml

@@ -13,6 +13,7 @@ on:
     - cron: "0 3 * * *"
 
 env:
+  PYBAMM_DISABLE_TELEMETRY: "true"
   FORCE_COLOR: 3
   PYBAMM_IDAKLU_EXPR_CASADI: ON
   PYBAMM_IDAKLU_EXPR_IREE: ON

.github/workflows/test_on_push.yml

@@ -5,6 +5,7 @@ on:
   pull_request:
 
 env:
+  PYBAMM_DISABLE_TELEMETRY: "true"
   FORCE_COLOR: 3
   PYBAMM_IDAKLU_EXPR_CASADI: ON
   PYBAMM_IDAKLU_EXPR_IREE: ON
@@ -36,7 +37,6 @@ jobs:
           pre-commit run -a
 
   run_unit_integration_and_coverage_tests:
-    needs: style
     runs-on: ${{ matrix.os }}
     strategy:
       fail-fast: false
@@ -132,7 +132,6 @@ jobs:
 
   # Skips IDAKLU module compilation for speedups, which is already tested in other jobs.
   run_doctests:
-    needs: style
     runs-on: ubuntu-latest
     strategy:
       fail-fast: false
@@ -177,7 +176,6 @@ jobs:
         run: python -m nox -s docs
 
   run_example_tests:
-    needs: style
     runs-on: ubuntu-latest
     strategy:
       fail-fast: false
@@ -233,7 +231,6 @@ jobs:
         run: python -m nox -s examples
 
   run_scripts_tests:
-    needs: style
     runs-on: ubuntu-latest
     strategy:
       fail-fast: false

.github/workflows/work_precision_sets.yml

@@ -5,6 +5,9 @@ on:
     types: [published]
   workflow_dispatch:
 
+env:
+  PYBAMM_DISABLE_TELEMETRY: "true"
+
 jobs:
   benchmarks_on_release:
     if: github.repository_owner == 'pybamm-team'

noxfile.py

@@ -49,7 +49,6 @@ def set_iree_state():
     "IREE_INDEX_URL": os.getenv(
         "IREE_INDEX_URL", "https://iree.dev/pip-release-links.html"
     ),
-    "PYBAMM_DISABLE_TELEMETRY": "true",
 }
 VENV_DIR = Path("./venv").resolve()
 

src/pybamm/config.py

@@ -8,6 +8,18 @@
 import time
 
 
+def check_env_opt_out():
+    return os.getenv("PYBAMM_DISABLE_TELEMETRY", "false").lower() != "false"
+
+
+def check_opt_out():
+    opt_out = check_env_opt_out()
+    config = pybamm.config.read()
+    if config:
+        opt_out = opt_out or not config["enable_telemetry"]
+    return opt_out
+
+
 def is_running_tests():  # pragma: no cover
     """
     Detect if the code is being run as part of a test suite or building docs with Sphinx.
@@ -21,20 +33,18 @@ def is_running_tests():  # pragma: no cover
     ):
         return True
 
-    # Check for GitHub Actions environment variable
-    if "GITHUB_ACTIONS" in os.environ:
-        return True
-
     # Check for other common CI environment variables
-    ci_env_vars = ["CI", "TRAVIS", "CIRCLECI", "JENKINS_URL", "GITLAB_CI"]
+    ci_env_vars = [
+        "GITHUB_ACTIONS",
+        "CI",
+        "TRAVIS",
+        "CIRCLECI",
+        "JENKINS_URL",
+        "GITLAB_CI",
+    ]
     if any(var in os.environ for var in ci_env_vars):
         return True
 
-    # Check for common test runner names in command-line arguments
-    test_runners = ["pytest", "unittest", "nose", "trial", "nox", "tox"]
-    if any(runner in arg.lower() for arg in sys.argv for runner in test_runners):
-        return True
-
     # Check if building docs with Sphinx
     if any(mod == "sphinx" or mod.startswith("sphinx.") for mod in sys.modules):
         print(
@@ -110,7 +120,7 @@ def get_input():  # pragma: no cover
 
 
 def generate():
-    if is_running_tests():
+    if is_running_tests() or check_opt_out():
         return
 
     # Check if the config file already exists

src/pybamm/telemetry.py

@@ -1,38 +1,45 @@
 from posthog import Posthog
-import os
 import pybamm
 import sys
 
-_posthog = Posthog(
-    # this is the public, write only API key, so it's ok to include it here
-    project_api_key="phc_bLZKBW03XjgiRhbWnPsnKPr0iw0z03fA6ZZYjxgW7ej",
-    host="https://us.i.posthog.com",
-)
 
-_posthog.log.setLevel("CRITICAL")
+class MockTelemetry:
+    def __init__(self):
+        self.disabled = True
 
+    @staticmethod
+    def capture(**kwargs):  # pragma: no cover
+        pass
 
-def disable():
-    _posthog.disabled = True
+
+if pybamm.config.check_opt_out():
+    _posthog = MockTelemetry()
+else:  # pragma: no cover
+    _posthog = Posthog(
+        # this is the public, write only API key, so it's ok to include it here
+        project_api_key="phc_bLZKBW03XjgiRhbWnPsnKPr0iw0z03fA6ZZYjxgW7ej",
+        host="https://us.i.posthog.com",
+    )
+    _posthog.log.setLevel("CRITICAL")
 
 
-_opt_out = os.getenv("PYBAMM_DISABLE_TELEMETRY", "false").lower()
-if _opt_out != "false":  # pragma: no cover
-    disable()
+def disable():
+    _posthog.disabled = True
 
 
 def capture(event):  # pragma: no cover
-    # don't capture events in automated testing
     if pybamm.config.is_running_tests() or _posthog.disabled:
         return
 
-    properties = {
-        "python_version": sys.version,
-        "pybamm_version": pybamm.__version__,
-    }
+    if pybamm.config.check_opt_out():
+        disable()
+        return
 
     config = pybamm.config.read()
     if config:
-        if config["enable_telemetry"]:
-            user_id = config["uuid"]
-            _posthog.capture(user_id, event, properties=properties)
+        properties = {
+            "python_version": sys.version,
+            "pybamm_version": pybamm.__version__,
+        }
+        user_id = config["uuid"]
+        _posthog.capture(distinct_id=user_id, event=event, properties=properties)

tests/shared.py

@@ -336,12 +336,12 @@ def no_internet_connection():
         conn = socket.create_connection((host, 80), 2)
         conn.close()
         return False
-    except socket.gaierror:
+    except (socket.gaierror, TimeoutError):
         return True
 
 
 def assert_domain_equal(a, b):
-    "Check that two domains are equal, ignoring empty domains"
+    """Check that two domains are equal, ignoring empty domains"""
     a_dict = {k: v for k, v in a.items() if v != []}
     b_dict = {k: v for k, v in b.items() if v != []}
     assert a_dict == b_dict

tests/unit/test_config.py

@@ -108,10 +108,8 @@ def mock_select(*args, **kwargs):
         assert "Timeout reached. Defaulting to not enabling telemetry." in captured.out
 
     def test_generate_and_read(self, monkeypatch, tmp_path):
-        # Mock is_running_tests to return False
         monkeypatch.setattr(pybamm.config, "is_running_tests", lambda: False)
-
-        # Mock ask_user_opt_in to return True
+        monkeypatch.setattr(pybamm.config, "check_opt_out", lambda: False)
         monkeypatch.setattr(pybamm.config, "ask_user_opt_in", lambda: True)
 
         # Mock telemetry capture
@@ -155,3 +153,17 @@ def test_read_uuid_from_file_invalid_yaml(self, tmp_path):
         config_dict = pybamm.config.read_uuid_from_file(invalid_yaml)
 
         assert config_dict is None
+
+    def test_check_opt_out(self, monkeypatch):
+        monkeypatch.setattr(pybamm.config, "read", lambda: {"enable_telemetry": True})
+        monkeypatch.setattr(pybamm.config, "check_env_opt_out", lambda: True)
+        assert pybamm.config.check_opt_out() is True
+        monkeypatch.setattr(pybamm.config, "read", lambda: {"enable_telemetry": False})
+        monkeypatch.setattr(pybamm.config, "check_env_opt_out", lambda: True)
+        assert pybamm.config.check_opt_out() is True
+        monkeypatch.setattr(pybamm.config, "read", lambda: {"enable_telemetry": True})
+        monkeypatch.setattr(pybamm.config, "check_env_opt_out", lambda: False)
+        assert pybamm.config.check_opt_out() is False
+        monkeypatch.setattr(pybamm.config, "read", lambda: {"enable_telemetry": False})
+        monkeypatch.setattr(pybamm.config, "check_env_opt_out", lambda: False)
+        assert pybamm.config.check_opt_out() is True