test: use freezegun instead of sleeps in time-sensitive tests (#957)

test: centralize freezegun mock time helpers and helper dependencies to testing package

Author: bndikt

diracx-db/tests/jobs/test_sandbox_metadata.py

@@ -1,8 +1,7 @@
 from __future__ import annotations
 
-import asyncio
 import secrets
-from datetime import datetime
+from datetime import datetime, timedelta
 
 import pytest
 import sqlalchemy
@@ -11,12 +10,14 @@
 from diracx.core.models import SandboxInfo, UserInfo
 from diracx.db.sql.sandbox_metadata.db import SandboxMetadataDB
 from diracx.db.sql.sandbox_metadata.schema import SandBoxes, SBEntityMapping
+from diracx.testing.time import install_sqlite_time_mock
 
 
 @pytest.fixture
 async def sandbox_metadata_db(tmp_path):
     sandbox_metadata_db = SandboxMetadataDB("sqlite+aiosqlite:///:memory:")
     async with sandbox_metadata_db.engine_context():
+        install_sqlite_time_mock(sandbox_metadata_db.engine)
         async with sandbox_metadata_db.engine.begin() as conn:
             await conn.run_sync(sandbox_metadata_db.metadata.create_all)
         yield sandbox_metadata_db
@@ -39,7 +40,7 @@ def test_get_pfn(sandbox_metadata_db: SandboxMetadataDB):
     )
 
 
-async def test_insert_sandbox(sandbox_metadata_db: SandboxMetadataDB):
+async def test_insert_sandbox(sandbox_metadata_db: SandboxMetadataDB, frozen_time):
     # TODO: DAL tests should be very simple, such complex tests should be handled in diracx-routers
     user_info = UserInfo(
         sub="vo:sub", preferred_username="user1", dirac_group="group1", vo="vo"
@@ -64,7 +65,7 @@ async def test_insert_sandbox(sandbox_metadata_db: SandboxMetadataDB):
     db_contents = await _dump_db(sandbox_metadata_db)
     owner_id1, last_access_time1 = db_contents[pfn1]
 
-    await asyncio.sleep(1)  # The timestamp only has second precision
+    frozen_time.tick(delta=timedelta(seconds=1))
     async with sandbox_metadata_db:
         with pytest.raises(SandboxAlreadyInsertedError):
             await sandbox_metadata_db.insert_sandbox(owner_id, "SandboxSE", pfn1, 100)
@@ -81,7 +82,7 @@ async def test_insert_sandbox(sandbox_metadata_db: SandboxMetadataDB):
         assert not await sandbox_metadata_db.sandbox_is_assigned(pfn1, "SandboxSE")
 
     # Inserting again should update the last access time
-    await asyncio.sleep(1)  # The timestamp only has second precision
+    frozen_time.tick(delta=timedelta(seconds=1))
     last_access_time3 = (await _dump_db(sandbox_metadata_db))[pfn1][1]
     assert last_access_time2 == last_access_time3
     async with sandbox_metadata_db:

diracx-db/tests/test_freeze_time.py

@@ -8,7 +8,7 @@
 import sqlalchemy
 from sqlalchemy.ext.asyncio import create_async_engine
 
-from diracx.testing.time import julian_date, mock_sqlite_time
+from diracx.testing.time import install_sqlite_time_mock, julian_date
 
 RE_SQLITE_TIME = re.compile(r"(\d{4})-(\d{2})-(\d{2})(?: (\d{2}):(\d{2}):(\d{2}))?")
 
@@ -40,7 +40,7 @@ async def test_freeze_sqlite_datetime(with_mock):
     """Test the SQLite DATETIME() function with freezegun."""
     engine = create_async_engine("sqlite+aiosqlite:///:memory:", future=True, echo=True)
     if with_mock:
-        sqlalchemy.event.listen(engine.sync_engine, "connect", mock_sqlite_time)
+        install_sqlite_time_mock(engine)
 
     async with engine.begin() as conn:
         # DATETIME()
@@ -83,7 +83,7 @@ async def test_freeze_sqlite_julianday(with_mock):
     """Test the SQLite JULIANDAY() function with freezegun."""
     engine = create_async_engine("sqlite+aiosqlite:///:memory:", future=True, echo=True)
     if with_mock:
-        sqlalchemy.event.listen(engine.sync_engine, "connect", mock_sqlite_time)
+        install_sqlite_time_mock(engine)
 
     async with engine.begin() as conn:
         # JULIANDAY()

diracx-logic/tests/jobs/test_sandboxes.py

@@ -10,7 +10,6 @@
 import httpx2
 import pytest
 import signurlarity.exceptions
-import sqlalchemy
 
 from diracx.core.exceptions import SandboxNotFoundError
 from diracx.core.models import ChecksumAlgorithm, SandboxFormat, SandboxInfo, UserInfo
@@ -21,7 +20,7 @@
     get_sandbox_file,
     initiate_sandbox_upload,
 )
-from diracx.testing.time import mock_sqlite_time
+from diracx.testing.time import install_sqlite_time_mock
 
 FAKE_USER_INFO = UserInfo(
     sub="fakevo:97ae90d3-36aa-4271-becf-e61173d93fe3",
@@ -36,7 +35,7 @@ async def sandbox_metadata_db() -> AsyncGenerator[SandboxMetadataDB, None]:
     """Create a fake sandbox metadata database."""
     db = SandboxMetadataDB(db_url="sqlite+aiosqlite:///:memory:")
     async with db.engine_context():
-        sqlalchemy.event.listen(db.engine.sync_engine, "connect", mock_sqlite_time)
+        install_sqlite_time_mock(db.engine)
 
         async with db.engine.begin() as conn:
             await conn.run_sync(db.metadata.create_all)

diracx-logic/tests/jobs/test_status.py

@@ -4,14 +4,13 @@
 from datetime import datetime, timezone
 
 import pytest
-import sqlalchemy
 
 from diracx.core.models import JobMetaData
 from diracx.db.os.job_parameters import JobParametersDB as RealJobParametersDB
 from diracx.db.sql.job.db import JobDB
 from diracx.logic.jobs import set_job_parameters_or_attributes
 from diracx.testing.mock_osdb import MockOSDBMixin
-from diracx.testing.time import mock_sqlite_time
+from diracx.testing.time import install_sqlite_time_mock
 
 
 # Reuse the generic MockOSDBMixin to build a mock JobParameters DB implementation
@@ -36,7 +35,7 @@ async def job_db() -> AsyncGenerator[JobDB, None]:
     """Create a fake sandbox metadata database."""
     db = JobDB(db_url="sqlite+aiosqlite:///:memory:")
     async with db.engine_context():
-        sqlalchemy.event.listen(db.engine.sync_engine, "connect", mock_sqlite_time)
+        install_sqlite_time_mock(db.engine)
 
         async with db.engine.begin() as conn:
             await conn.run_sync(db.metadata.create_all)

diracx-routers/tests/auth/test_standard.py

@@ -12,7 +12,6 @@
 import jwt
 import pytest
 from cryptography.fernet import Fernet
-from freezegun import freeze_time
 from joserfc.jwk import RSAKey, OKPKey, KeySet
 from joserfc.errors import (
     UnsupportedKeyOperationError,
@@ -562,56 +561,60 @@ async def test_refresh_token_rotation(test_client, auth_httpx_mock: HTTPXMock):
 
 
 async def test_refresh_token_expired(
-    test_client, test_auth_settings: AuthSettings, auth_httpx_mock: HTTPXMock
+    test_client,
+    test_auth_settings: AuthSettings,
+    auth_httpx_mock: HTTPXMock,
+    frozen_time,
 ):
     """Test the expiration date of the passed refresh token.
 
     - get a refresh token
     - move time forward past its expiration time
     - ensure the expired token is rejected.
     """
-    with freeze_time(datetime.now(tz=timezone.utc)) as frozen_time:
-        # Get refresh token
-        refresh_token = _get_tokens(test_client)["refresh_token"]
+    # Get refresh token
+    refresh_token = _get_tokens(test_client)["refresh_token"]
 
-        frozen_time.tick(
-            delta=timedelta(minutes=test_auth_settings.refresh_token_expire_minutes + 1)
-        )
+    frozen_time.tick(
+        delta=timedelta(minutes=test_auth_settings.refresh_token_expire_minutes + 1)
+    )
 
-        request_data = {
-            "grant_type": "refresh_token",
-            "refresh_token": refresh_token,
-            "client_id": DIRAC_CLIENT_ID,
-        }
+    request_data = {
+        "grant_type": "refresh_token",
+        "refresh_token": refresh_token,
+        "client_id": DIRAC_CLIENT_ID,
+    }
 
-        # Try to get a new access token using the expired refresh token
-        r = test_client.post("/api/auth/token", data=request_data)
+    # Try to get a new access token using the expired refresh token
+    r = test_client.post("/api/auth/token", data=request_data)
     data = r.json()
     assert r.status_code == 401, data
     assert data["detail"] == "expired_token: The token is expired"
 
 
 async def test_access_token_expired(
-    test_client, test_auth_settings: AuthSettings, auth_httpx_mock: HTTPXMock
+    test_client,
+    test_auth_settings: AuthSettings,
+    auth_httpx_mock: HTTPXMock,
+    frozen_time,
 ):
     """Test the expiration date of the passed access token.
 
     - get an access token
     - move time forward past its expiration time
     - ensure the expired token is rejected.
     """
-    with freeze_time(datetime.now(tz=timezone.utc)) as frozen_time:
-        # Get access token
-        access_token = _get_tokens(test_client)["access_token"]
+    # Get access token
+    access_token = _get_tokens(test_client)["access_token"]
 
-        frozen_time.tick(
-            delta=timedelta(minutes=test_auth_settings.access_token_expire_minutes + 1)
-        )
+    frozen_time.tick(
+        delta=timedelta(minutes=test_auth_settings.access_token_expire_minutes + 1)
+    )
 
-        headers = {"Authorization": f"Bearer {access_token}"}
+    headers = {"Authorization": f"Bearer {access_token}"}
 
-        # Try to get the userinfo using the expired access token
-        r = test_client.get("/api/auth/userinfo", headers=headers)
+    # Try to get the userinfo using the expired access token
+    r = test_client.get("/api/auth/userinfo", headers=headers)
     data = r.json()
     assert r.status_code == 401, data
     assert data["detail"] == "Invalid JWT"

diracx-routers/tests/jobs/test_heartbeat_commands.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 from datetime import datetime, timedelta, timezone
-from time import sleep
 
 import pytest
 from fastapi.testclient import TestClient
@@ -23,7 +22,7 @@
 )
 
 
-def test_heartbeat(normal_user_client: TestClient, valid_job_id: int):
+def test_heartbeat(frozen_time, normal_user_client: TestClient, valid_job_id: int):
     search_body = {
         "search": [{"parameter": "JobID", "operator": "eq", "value": valid_job_id}]
     }
@@ -60,7 +59,7 @@ def test_heartbeat(normal_user_client: TestClient, valid_job_id: int):
     )
     r.raise_for_status()
 
-    sleep(1)
+    frozen_time.tick(delta=timedelta(seconds=1))
     # Send another heartbeat and check that a Kill job command was set
     payload = {valid_job_id: {"Vsize": 1235}}
     r = normal_user_client.patch("/api/jobs/heartbeat", json=payload)
@@ -74,7 +73,7 @@ def test_heartbeat(normal_user_client: TestClient, valid_job_id: int):
     assert commands[0]["command"] == "Kill", (
         f"Wrong job command received, should be 'Kill' but got {commands[0]=}"
     )
-    sleep(1)
+    frozen_time.tick(delta=timedelta(seconds=1))
 
     # Send another heartbeat and check the job commands are empty
     payload = {valid_job_id: {"Vsize": 1234}}
@@ -87,6 +86,7 @@ def test_heartbeat(normal_user_client: TestClient, valid_job_id: int):
 
 
 def test_multiple_jobs_receive_independent_kill_commands(
+    frozen_time,
     normal_user_client: TestClient,
     valid_job_ids: list[int],
 ):
@@ -105,7 +105,7 @@ def test_multiple_jobs_receive_independent_kill_commands(
     )
     r.raise_for_status()
 
-    sleep(1)
+    frozen_time.tick(delta=timedelta(seconds=1))
 
     r = normal_user_client.patch(
         "/api/jobs/heartbeat",
@@ -118,7 +118,7 @@ def test_multiple_jobs_receive_independent_kill_commands(
     assert {cmd["job_id"] for cmd in commands} == set(valid_job_ids)
     assert {cmd["command"] for cmd in commands} == {"Kill"}
 
-    sleep(1)
+    frozen_time.tick(delta=timedelta(seconds=1))
 
     r = normal_user_client.patch(
         "/api/jobs/heartbeat",
@@ -130,6 +130,7 @@ def test_multiple_jobs_receive_independent_kill_commands(
 
 
 def test_non_killed_status_does_not_create_command(
+    frozen_time,
     normal_user_client: TestClient,
     valid_job_id: int,
 ):
@@ -147,7 +148,7 @@ def test_non_killed_status_does_not_create_command(
     )
     r.raise_for_status()
 
-    sleep(1)
+    frozen_time.tick(delta=timedelta(seconds=1))
 
     r = normal_user_client.patch(
         "/api/jobs/heartbeat",
@@ -159,6 +160,7 @@ def test_non_killed_status_does_not_create_command(
 
 
 def test_deleted_creates_kill_command(
+    frozen_time,
     normal_user_client: TestClient,
     valid_job_id: int,
 ):
@@ -176,7 +178,7 @@ def test_deleted_creates_kill_command(
     )
     r.raise_for_status()
 
-    sleep(1)
+    frozen_time.tick(delta=timedelta(seconds=1))
 
     r = normal_user_client.patch(
         "/api/jobs/heartbeat",
@@ -189,7 +191,7 @@ def test_deleted_creates_kill_command(
     assert commands[0]["job_id"] == valid_job_id
     assert commands[0]["command"] == "Kill"
 
-    sleep(1)
+    frozen_time.tick(delta=timedelta(seconds=1))
 
     r = normal_user_client.patch(
         "/api/jobs/heartbeat",

diracx-testing/pyproject.toml

@@ -20,6 +20,8 @@ dependencies = [
     "pytest-xdist",
     "httpx2",
     "joserfc",
+    "freezegun",
+    "sqlalchemy",
     "uuid-utils",
     "pytest-github-actions-annotate-failures",
 ]

diracx-testing/src/diracx/testing/__init__.py

@@ -10,6 +10,7 @@
     "demo_urls",
     "do_device_flow_with_dex",
     "fernet_key",
+    "frozen_time",
     "private_key",
     "pytest_addoption",
     "session_client_factory",
@@ -23,6 +24,7 @@
 ]
 
 from .entrypoints import verify_entry_points
+from .time import frozen_time
 from .utils import (
     ClientFactory,
     aio_moto,