chore(tests): add test suite for semaphore usage (e.g. multithreaded env) (#2596)

Author: haakonvt

cognite/client/utils/_concurrency.py

@@ -262,6 +262,13 @@ def __init__(self) -> None:
             write_schema=1,
         )
 
+    @functools.cached_property
+    def _all_concurrency_configs(self) -> list[ConcurrencyConfig]:
+        """Helper method primarily used in testing to handle the 'annoying' state of concurrency settings"""
+        configs = [name for name, val in vars(type(self)).items() if isinstance(val, property)]
+        configs.remove("is_frozen")
+        return [getattr(self, name) for name in configs]
+
     def _check_frozen(self, name: str, api_name: str) -> None:
         if self.__frozen:
             raise RuntimeError(
@@ -556,22 +563,28 @@ def start(self) -> None:
 
 
 # We need this in order to support a synchronous Cognite client.
-_INTERNAL_EVENT_LOOP_THREAD_EXECUTOR_SINGLETON: EventLoopThreadExecutor
+_INTERNAL_EVENT_LOOP_THREAD_EXECUTOR_SINGLETON: EventLoopThreadExecutor | None = None
+_EXECUTOR_INIT_LOCK = threading.Lock()
 
 
 def _get_event_loop_executor() -> EventLoopThreadExecutor:
     global _INTERNAL_EVENT_LOOP_THREAD_EXECUTOR_SINGLETON
-    try:
+
+    # Fast path: singleton already initialized — no lock needed.
+    if _INTERNAL_EVENT_LOOP_THREAD_EXECUTOR_SINGLETON is not None:
+        return _INTERNAL_EVENT_LOOP_THREAD_EXECUTOR_SINGLETON
+
+    # Slow path: serialize initialization. Without this, multiple threads racing on the first
+    # call (e.g. a ThreadPoolExecutor of sync clients) can each construct their own executor
+    # and end up using different background loops, breaking the per-loop semaphore cache key.
+    with _EXECUTOR_INIT_LOCK:
+        if _INTERNAL_EVENT_LOOP_THREAD_EXECUTOR_SINGLETON is None:
+            ex_cls = EventLoopThreadExecutor
+            if _RUNNING_IN_BROWSER:
+                ex_cls = cast(type[EventLoopThreadExecutor], _PyodideEventLoopExecutor)
+            _INTERNAL_EVENT_LOOP_THREAD_EXECUTOR_SINGLETON = ex_cls()
+            _INTERNAL_EVENT_LOOP_THREAD_EXECUTOR_SINGLETON.start()
         return _INTERNAL_EVENT_LOOP_THREAD_EXECUTOR_SINGLETON
-    except NameError:
-        # First time we need to initialize:
-        ex_cls = EventLoopThreadExecutor
-        if _RUNNING_IN_BROWSER:
-            ex_cls = cast(type[EventLoopThreadExecutor], _PyodideEventLoopExecutor)
-
-        executor = _INTERNAL_EVENT_LOOP_THREAD_EXECUTOR_SINGLETON = ex_cls()
-        executor.start()
-        return executor
 
 
 async def execute_async_tasks_with_fail_fast(tasks: list[AsyncSDKTask]) -> TasksSummary:

pytest.ini

@@ -1,7 +1,8 @@
 [pytest]
 markers =
-    dsl
-    coredeps
+    dsl: requires optional data-science deps (e.g. pandas/numpy) - skipped when only core deps are installed
+    coredeps: runs only with --test-deps-only-core (verifies that optional deps are NOT installed in the core-only env)
+    allow_no_semaphore: opts a test out of the strict semaphore check (see tests/conftest.py)
 
 anyio_mode = auto
 

tests/conftest.py

@@ -2,18 +2,56 @@
 
 import platform
 from collections.abc import Callable, Iterator
+from typing import Any
 
 import dotenv
 import pytest
 from _pytest.monkeypatch import MonkeyPatch
 
 from cognite.client import global_config
 from cognite.client._api_client import APIClient
+from cognite.client._http_client import AsyncHTTPClientWithRetry
 
 dotenv.load_dotenv()
 
 global_config.disable_pypi_version_check = True
 
+ALLOW_NO_SEMAPHORE_MARKER = "allow_no_semaphore"
+
+
+@pytest.fixture(autouse=True)
+def require_semaphore_on_every_request(request: pytest.FixtureRequest, monkeypatch: pytest.MonkeyPatch) -> None:
+    """As a main rule, all API calls in the test suite MUST route through a semaphore.
+
+    Production code falls back to ``nullcontext()`` when ``semaphore is None``; that path
+    exists for users calling top-level ``client.post(...)``/``client.get(...)`` with raw URLs,
+    but no internal API method should ever hit it. Patching here turns a missing semaphore
+    into a hard failure so a regression (forgotten ``semaphore=...`` arg) shows up loudly.
+
+    Tests that legitimately exercise the None path opt out via ``@pytest.mark.allow_no_semaphore("<reason>")``.
+    These are the top-level methods mentioned above plus any method entering the semaphore
+    at a higher level than at the HTTP request level (e.g. datapoints.insert)
+    """
+    if request.node.get_closest_marker(ALLOW_NO_SEMAPHORE_MARKER):
+        return
+
+    original = AsyncHTTPClientWithRetry._with_retry
+
+    async def strict(
+        self: AsyncHTTPClientWithRetry, coro_factory: Any, *, url: str, headers: Any, semaphore: Any
+    ) -> Any:
+        if semaphore is None:
+            pytest.fail(
+                f"Internal API call to {url!r} did not pass a semaphore. "
+                "All endpoints behind client.<api>.<method> must route through a semaphore — "
+                "the nullcontext fallback is reserved for top-level client.post/get calls only. "
+                "If this call legitimately holds the semaphore at a higher level, mark the test "
+                f"with @pytest.mark.{ALLOW_NO_SEMAPHORE_MARKER}('<reason>')."
+            )
+        return await original(self, coro_factory, url=url, headers=headers, semaphore=semaphore)
+
+    monkeypatch.setattr(AsyncHTTPClientWithRetry, "_with_retry", strict)
+
 
 @pytest.fixture(scope="session")
 def anyio_backend() -> str:

tests/tests_integration/test_api/test_datapoint_subscriptions.py

@@ -2,18 +2,16 @@
 
 import math
 import random
-import time
 import unittest
 from collections.abc import Iterator
 from contextlib import contextmanager
-from datetime import datetime
 
 import numpy as np
 import pandas as pd
 import pytest
 
 from cognite.client import CogniteClient
-from cognite.client.data_classes import TimeSeries, TimeSeriesWrite, filters
+from cognite.client.data_classes import TimeSeriesWrite, filters
 from cognite.client.data_classes.data_modeling import NodeId, SpaceApply
 from cognite.client.data_classes.data_modeling.cdm.v1 import CogniteTimeSeriesApply
 from cognite.client.data_classes.datapoints_subscriptions import (
@@ -302,6 +300,7 @@ def test_iterate_data_subscription_changed_time_series(
             assert {a.external_id for a in batch.subscription_changes.removed} == {first_ts}
             assert len(batch.updates) == 0
 
+    @pytest.mark.allow_no_semaphore("Part of the tests hits DatapointsAPI._insert_datapoints")
     def test_iterate_data_subscription_datapoints_added(
         self, cognite_client: CogniteClient, time_series_external_ids: list[str]
     ) -> None:
@@ -384,6 +383,7 @@ def test_iterate_data_subscription_start_1m_ago(
             "you run these tests, please wait 1 minute and try again.",
         )
 
+    @pytest.mark.allow_no_semaphore("Part of the tests hits DatapointsAPI._insert_datapoints")
     def test_iterate_data__using_status_codes(
         self, cognite_client: CogniteClient, sub_for_status_codes: DatapointSubscription
     ) -> None:
@@ -429,63 +429,3 @@ def test_iterate_data__using_status_codes(
         assert bad_upsert_value[0] is None
         assert all(isinstance(v, float) for v in bad_upsert_value[1:])
         assert no_bad[0].upserts.status_symbol == ["Uncertain", "Good", "Good", "Good"]
-
-    @pytest.mark.skip(reason="Using a filter (as opposed to specific identifiers) is eventually consistent")
-    def test_update_filter_subscription_added_times_series(
-        self, cognite_client: CogniteClient, time_series_external_ids: list[str]
-    ) -> None:
-        f = filters
-        p = DatapointSubscriptionProperty
-        numerical_timeseries = f.And(
-            f.Equals(p.is_string, False), f.Prefix(p.external_id, "PYSDK DataPoint Subscription Test")
-        )
-        new_subscription = DataPointSubscriptionWrite(
-            external_id=f"PYSDKDataPointSubscriptionUpdateFilterTest-{random_string(10)}",
-            name="PYSDKDataPointSubscriptionUpdateFilterTest",
-            filter=numerical_timeseries,
-            partition_count=1,
-        )
-        created_timeseries: TimeSeries | None = None
-        with create_subscription_with_cleanup(cognite_client, new_subscription) as created:
-            assert created.created_time
-
-            initial_added_count = 0
-            for batch in cognite_client.time_series.subscriptions.iterate_data(
-                new_subscription.external_id,
-                poll_timeout=0,
-            ):
-                initial_added_count += len(batch.subscription_changes.added)
-                if not batch.has_next:
-                    break
-
-            assert initial_added_count > 0, "There should be at least one numerical timeseries added"
-
-            new_numerical_timeseries = TimeSeriesWrite(
-                external_id=f"PYSDK DataPoint Subscription Test 42 ({random_string(10)})",
-                name="PYSDK DataPoint Subscription Test 42",
-                is_string=False,
-            )
-            try:
-                created_timeseries = cognite_client.time_series.create(new_numerical_timeseries)
-                cognite_client.time_series.data.insert(
-                    [(datetime.now(), 42)], external_id=new_numerical_timeseries.external_id
-                )
-                # Ensure that the subscription has been updated
-                time.sleep(10)
-
-                updated_added_count = 0
-                for batch in cognite_client.time_series.subscriptions.iterate_data(
-                    new_subscription.external_id,
-                    poll_timeout=0,
-                ):
-                    updated_added_count += len(batch.subscription_changes.added)
-                    if not batch.has_next:
-                        break
-
-                assert initial_added_count + 1 == updated_added_count, (
-                    "The new timeseries should be added. This is most likely because using a filter with "
-                    "datapoint subscriptions is eventually consistent."
-                )
-            finally:
-                if created_timeseries:
-                    cognite_client.time_series.delete(created_timeseries.id)

tests/tests_integration/test_api/test_datapoints.py

@@ -501,6 +501,10 @@ def ts_create_in_dms(
 
 
 class TestTimeSeriesCreatedInDMS:
+    @pytest.mark.allow_no_semaphore(
+        "Test inserts datapoints; DatapointsAPI._insert_datapoints holds the semaphore via outer "
+        "'async with' and passes None to _post to avoid double-acquiring."
+    )
     def test_insert_read_delete_dps(self, cognite_client: CogniteClient, ts_create_in_dms: NodeApplyResult) -> None:
         # Ensure the DMS time series is retrievable from normal TS API:
         inst_id = ts_create_in_dms.as_id()
@@ -2855,6 +2859,11 @@ def test_instance_id_and_missing(self, cognite_client: CogniteClient, instance_t
             cognite_client.time_series.data.retrieve_latest(instance_id=missing, ignore_unknown_ids=False)
 
 
+@pytest.mark.allow_no_semaphore(
+    "Insert paths go through DatapointsAPI._insert_datapoints, which holds the semaphore via "
+    "outer 'async with' for memory backpressure and passes None to _post. Delete paths in this "
+    "class use the regular semaphore route and are still checked against the broader suite."
+)
 class TestInsertDatapointsAPI:
     @pytest.mark.usefixtures("post_spy")
     def test_insert(

tests/tests_integration/test_api/test_event_loop_interop.py

@@ -6,26 +6,23 @@
 from __future__ import annotations
 
 import random
+from collections.abc import Iterator
 from concurrent.futures import ThreadPoolExecutor
 from typing import Any
 
 import pytest
 
-from cognite.client import AsyncCogniteClient, CogniteClient, global_config
+from cognite.client import AsyncCogniteClient, CogniteClient
 from cognite.client._http_client import _global_async_httpx_clients
+from tests.utils import fresh_concurrency_state
 
 
 @pytest.fixture(autouse=True)
-def _clear_loop_caches() -> None:
+def _clear_loop_caches() -> Iterator[None]:
     """Reset per-loop caches so each test starts with a clean slate (...and state lol)."""
     _global_async_httpx_clients.clear()
-    for config in (
-        global_config.concurrency_settings.general,
-        global_config.concurrency_settings.datapoints,
-        global_config.concurrency_settings.raw,
-        global_config.concurrency_settings.data_modeling,
-    ):
-        config._semaphore_cache.clear()
+    with fresh_concurrency_state():
+        yield
 
 
 def make_sync_api_call(client: CogniteClient) -> Any:

tests/tests_integration/test_api/test_simulators/test_models.py

@@ -31,6 +31,11 @@
     "seed_resource_names",
     "seed_simulator_model_revisions",
 )
+@pytest.mark.allow_no_semaphore(
+    "Simulator tests use direct _post(..., semaphore=None) calls in fixtures/seed helpers to "
+    "exercise endpoints (logs/update, callback, integrations/update) that aren't exposed via "
+    "public SDK methods — they're worker-facing endpoints used here for test setup only."
+)
 class TestSimulatorModels:
     def test_list_models(self, cognite_client: CogniteClient, seed_resource_names: ResourceNames) -> None:
         models = cognite_client.simulators.models.list(

tests/tests_integration/test_api/test_simulators/test_runs.py

@@ -20,6 +20,10 @@
 from tests.tests_integration.test_api.test_simulators.seed.data import RESOURCES, ResourceNames
 
 
+@pytest.mark.allow_no_semaphore(
+    "Simulator tests use direct _post(..., semaphore=None) calls (e.g. /simulators/run/callback) "
+    "to exercise worker-facing endpoints not exposed via public SDK methods."
+)
 @pytest.mark.usefixtures("seed_resource_names", "seed_simulator_routine_revisions")
 class TestSimulatorRuns:
     def test_list_filtering(

tests/tests_integration/test_api/test_workflows.py

@@ -463,6 +463,10 @@ def test_list_workflows(self, cognite_client: CogniteClient, persisted_workflow_
 
 
 class TestWorkflowVersions:
+    @pytest.mark.allow_no_semaphore(
+        "Test setup uses simulator fixtures that exercise worker-facing endpoints "
+        "(/simulators/integrations/update) via direct _post(..., semaphore=None)."
+    )
     def test_upsert_run_delete_with_simulation_task(
         self,
         cognite_client: CogniteClient,

tests/tests_integration/test_cognite_client.py

@@ -19,6 +19,10 @@ def cognite_client_with_wrong_base_url(
     return cognite_client
 
 
+@pytest.mark.allow_no_semaphore(
+    "Tests exercise top-level client.post/get/put/token-inspect on the (Async)CogniteClient — "
+    "the documented nullcontext path. Concurrency limiting is unrelated."
+)
 class TestCogniteClient:
     def test_wrong_project(
         self, monkeypatch: MonkeyPatch, async_client: AsyncCogniteClient, cognite_client: CogniteClient