diff --git a/pyproject.toml b/pyproject.toml index 4d7ee91575..d6b0d5f07b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -193,7 +193,7 @@ test = [ "pytest-unordered>=0.7.0", "pytest-xdist>=3.6.1", "requests-mock>=1.12.1", - "freezegun>=1.5.1", + "time-machine>=3.5.1", "httpbin>=0.10.2", # indirect to make compatible with Werkzeug 3 "openai>=1.68.2", "testcontainers>=4.15.0", diff --git a/src/dstack/_internal/proxy/gateway/services/registry.py b/src/dstack/_internal/proxy/gateway/services/registry.py index 2ce9ea9ffd..e0f67868bb 100644 --- a/src/dstack/_internal/proxy/gateway/services/registry.py +++ b/src/dstack/_internal/proxy/gateway/services/registry.py @@ -1,6 +1,5 @@ import asyncio from asyncio import Lock -from datetime import datetime from pathlib import Path from typing import Iterable, Optional @@ -26,6 +25,7 @@ ServiceConnection, ServiceConnectionPool, ) +from dstack._internal.utils.common import get_current_datetime from dstack._internal.utils.logging import get_logger ACCESS_LOG_PATH = Path("/var/log/nginx/dstack.access.log") @@ -97,7 +97,7 @@ async def register_service( models.ChatModel( project_name=project_name, name=model.name, - created_at=datetime.now(), + created_at=get_current_datetime(), run_name=run_name, format_spec=model_schema_to_format_spec(model), ), diff --git a/src/tests/_internal/cli/utils/test_run.py b/src/tests/_internal/cli/utils/test_run.py index 35bf28a430..dffde505e8 100644 --- a/src/tests/_internal/cli/utils/test_run.py +++ b/src/tests/_internal/cli/utils/test_run.py @@ -4,7 +4,7 @@ from unittest.mock import Mock import pytest -from freezegun import freeze_time +import time_machine from rich.table import Table from rich.text import Text from sqlalchemy import select @@ -206,7 +206,7 @@ async def create_run_with_job( @pytest.mark.usefixtures("test_db", "image_config_mock") @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) class TestGetRunsTable: - @freeze_time(datetime(2026, 1, 1)) + @time_machine.travel(datetime(2026, 1, 1), tick=False) async def test_simple_run(self, session: AsyncSession): api_run = await create_run_with_job( session=session, submitted_at=datetime(2023, 1, 1, tzinfo=timezone.utc) diff --git a/src/tests/_internal/proxy/gateway/routers/test_registry.py b/src/tests/_internal/proxy/gateway/routers/test_registry.py index 393b901acc..a9b265d9fc 100644 --- a/src/tests/_internal/proxy/gateway/routers/test_registry.py +++ b/src/tests/_internal/proxy/gateway/routers/test_registry.py @@ -1,12 +1,12 @@ import re import uuid -from datetime import datetime +from datetime import datetime, timezone from pathlib import Path from typing import Optional import httpx import pytest -from freezegun import freeze_time +import time_machine from dstack._internal.core.errors import SSHError from dstack._internal.proxy.gateway.app import make_app @@ -209,7 +209,7 @@ async def test_register_same_domain_error(self, tmp_path: Path, system_mocks: Mo assert (tmp_path / "sites-enabled" / "443-test-run.gtw.test.conf").exists() assert system_mocks.reload_nginx.call_count == 1 - @freeze_time(datetime(2024, 12, 12, 0, 30)) + @time_machine.travel(datetime(2024, 12, 12, 0, 30, tzinfo=timezone.utc), tick=False) async def test_register_with_model(self, tmp_path: Path, system_mocks: Mocks) -> None: repo = GatewayProxyRepo() client = make_client(tmp_path, repo=repo) @@ -225,7 +225,7 @@ async def test_register_with_model(self, tmp_path: Path, system_mocks: Mocks) -> ChatModel( project_name="test-proj", name="test-model", - created_at=datetime(2024, 12, 12, 0, 30), + created_at=datetime(2024, 12, 12, 0, 30, tzinfo=timezone.utc), run_name="test-run", format_spec=OpenAIChatModelFormat(prefix="/v1"), ) diff --git a/src/tests/_internal/proxy/gateway/services/test_stats.py b/src/tests/_internal/proxy/gateway/services/test_stats.py index e019aed7f2..f949178a4b 100644 --- a/src/tests/_internal/proxy/gateway/services/test_stats.py +++ b/src/tests/_internal/proxy/gateway/services/test_stats.py @@ -3,14 +3,14 @@ from textwrap import dedent import pytest -from freezegun import freeze_time +import time_machine from dstack._internal.proxy.gateway.schemas.stats import Stat from dstack._internal.proxy.gateway.services.stats import StatsCollector @pytest.mark.asyncio -@freeze_time(datetime(2024, 12, 6, 12, 10, tzinfo=timezone.utc)) +@time_machine.travel(datetime(2024, 12, 6, 12, 10, tzinfo=timezone.utc), tick=False) @pytest.mark.parametrize( ("access_log", "expected_result"), [ @@ -128,7 +128,7 @@ async def test_collect_stats(access_log: str, expected_result: dict, tmp_path: P @pytest.mark.asyncio -@freeze_time(datetime(2024, 12, 6, 12, 10, tzinfo=timezone.utc)) +@time_machine.travel(datetime(2024, 12, 6, 12, 10, tzinfo=timezone.utc), tick=False) async def test_collect_stats_after_log_update(tmp_path: Path) -> None: access_log_path = tmp_path / "dstack.access.log" collector = StatsCollector(access_log_path) diff --git a/src/tests/_internal/server/background/pipeline_tasks/test_instances/test_termination.py b/src/tests/_internal/server/background/pipeline_tasks/test_instances/test_termination.py index b9da58fc11..3077fb8e75 100644 --- a/src/tests/_internal/server/background/pipeline_tasks/test_instances/test_termination.py +++ b/src/tests/_internal/server/background/pipeline_tasks/test_instances/test_termination.py @@ -4,7 +4,7 @@ from unittest.mock import AsyncMock, Mock, patch import pytest -from freezegun import freeze_time +import time_machine from sqlalchemy.ext.asyncio import AsyncSession from dstack._internal.core.errors import BackendError, NotYetTerminated @@ -126,7 +126,7 @@ async def test_terminate_retry( await session.commit() with ( - freeze_time(initial_time + dt.timedelta(minutes=1)), + time_machine.travel(initial_time + dt.timedelta(minutes=1), tick=False), self.mock_terminate_in_backend(error=error) as mock, ): await process_instance(session, worker, instance) @@ -135,7 +135,7 @@ async def test_terminate_retry( assert instance.status == InstanceStatus.TERMINATING with ( - freeze_time(initial_time + dt.timedelta(minutes=2)), + time_machine.travel(initial_time + dt.timedelta(minutes=2), tick=False), self.mock_terminate_in_backend(error=None) as mock, ): await process_instance(session, worker, instance) @@ -162,7 +162,7 @@ async def test_terminate_not_retries_if_too_early( await session.commit() with ( - freeze_time(initial_time + dt.timedelta(minutes=1)), + time_machine.travel(initial_time + dt.timedelta(minutes=1), tick=False), self.mock_terminate_in_backend(error=BackendError("err")) as mock, ): await process_instance(session, worker, instance) @@ -174,7 +174,7 @@ async def test_terminate_not_retries_if_too_early( await session.commit() with ( - freeze_time(initial_time + dt.timedelta(minutes=1, seconds=11)), + time_machine.travel(initial_time + dt.timedelta(minutes=1, seconds=11), tick=False), self.mock_terminate_in_backend(error=None) as mock, ): await process_instance(session, worker, instance) @@ -201,7 +201,7 @@ async def test_terminate_on_termination_deadline( await session.commit() with ( - freeze_time(initial_time + dt.timedelta(minutes=1)), + time_machine.travel(initial_time + dt.timedelta(minutes=1), tick=False), self.mock_terminate_in_backend(error=BackendError("err")) as mock, ): await process_instance(session, worker, instance) @@ -210,7 +210,7 @@ async def test_terminate_on_termination_deadline( assert instance.status == InstanceStatus.TERMINATING with ( - freeze_time(initial_time + dt.timedelta(minutes=15, seconds=55)), + time_machine.travel(initial_time + dt.timedelta(minutes=15, seconds=55), tick=False), self.mock_terminate_in_backend(error=None) as mock, ): await process_instance(session, worker, instance) diff --git a/src/tests/_internal/server/background/pipeline_tasks/test_running_jobs.py b/src/tests/_internal/server/background/pipeline_tasks/test_running_jobs.py index 1e13956ca7..6b82e0c3f0 100644 --- a/src/tests/_internal/server/background/pipeline_tasks/test_running_jobs.py +++ b/src/tests/_internal/server/background/pipeline_tasks/test_running_jobs.py @@ -8,7 +8,7 @@ from unittest.mock import AsyncMock, MagicMock, Mock, patch import pytest -from freezegun import freeze_time +import time_machine from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.orm import selectinload @@ -1281,7 +1281,7 @@ async def test_pulling_shim_failed( with ( patch("dstack._internal.server.services.runner.pool.SSHTunnel") as ssh_tunnel_cls, - freeze_time(job.disconnected_at + timedelta(minutes=5)), + time_machine.travel(job.disconnected_at + timedelta(minutes=5), tick=False), ): from dstack._internal.core.errors import SSHError @@ -1622,7 +1622,7 @@ async def test_running_job_disconnect_retries_then_terminates( with ( patch("dstack._internal.server.services.runner.pool.SSHTunnel") as ssh_tunnel_cls, - freeze_time(job.disconnected_at + timedelta(minutes=5)), + time_machine.travel(job.disconnected_at + timedelta(minutes=5), tick=False), ): ssh_tunnel_cls.side_effect = SSHError await _process_job(session, worker, job) @@ -1858,7 +1858,7 @@ async def test_max_duration_enforced_by_server( patch( "dstack._internal.server.services.runner.client.RunnerClient.from_address" ) as runner_client_cls, - freeze_time(now), + time_machine.travel(now, tick=False), ): runner_client_mock = runner_client_cls.return_value runner_client_mock.pull.return_value = PullResponse( @@ -1909,7 +1909,7 @@ async def test_max_duration_enforced_by_server( ), ], ) - @freeze_time(datetime(2023, 1, 1, 12, 30, tzinfo=timezone.utc)) + @time_machine.travel(datetime(2023, 1, 1, 12, 30, tzinfo=timezone.utc), tick=False) async def test_gpu_utilization( self, test_db, @@ -3148,7 +3148,7 @@ async def test_router_failed_terminates_worker(self): result.job_update_map.get("termination_reason_message") or "" ) - @freeze_time("2023-01-01 12:00:00Z") + @time_machine.travel(datetime(2023, 1, 1, 12, tzinfo=timezone.utc), tick=False) async def test_router_not_provisioned_within_timeout_defers(self): context = self._make_context( submitted_at=datetime(2023, 1, 1, 11, 45, 0, tzinfo=timezone.utc), @@ -3162,7 +3162,7 @@ async def test_router_not_provisioned_within_timeout_defers(self): assert out is None assert result.job_update_map == {} - @freeze_time("2023-01-01 12:00:00Z") + @time_machine.travel(datetime(2023, 1, 1, 12, tzinfo=timezone.utc), tick=False) async def test_router_not_provisioned_past_timeout_terminates(self): context = self._make_context( submitted_at=datetime(2023, 1, 1, 10, 0, 0, tzinfo=timezone.utc), @@ -3262,7 +3262,7 @@ def _make_context(self) -> _ProcessContext: instance_access_revoked=False, ) - @freeze_time("2023-01-01 12:00:00Z") + @time_machine.travel(datetime(2023, 1, 1, 12, tzinfo=timezone.utc), tick=False) async def test_submitted_sibling_defers(self): context = self._make_context() result = _ProcessResult() @@ -3345,7 +3345,7 @@ async def _fake_session_ctx(): ), ) - @freeze_time("2023-01-01 12:00:00Z") + @time_machine.travel(datetime(2023, 1, 1, 12, tzinfo=timezone.utc), tick=False) async def test_groups_ip_not_ready_defers(self): context = self._make_context() result = _ProcessResult() @@ -3556,7 +3556,7 @@ async def _fake_session_ctx(): ), ) - @freeze_time("2023-01-01 12:00:00Z") + @time_machine.travel(datetime(2023, 1, 1, 12, tzinfo=timezone.utc), tick=False) async def test_replica_ip_not_ready_defers(self): context = self._make_context(router_ip="") result = _ProcessResult() @@ -3568,7 +3568,7 @@ async def test_replica_ip_not_ready_defers(self): assert result.job_update_map == {} assert context.job.job_spec.commands == ["echo ${{ groups[0].replicas[0].IP_ADDRESS }}"] - @freeze_time("2023-01-01 12:00:00Z") + @time_machine.travel(datetime(2023, 1, 1, 12, tzinfo=timezone.utc), tick=False) async def test_replica_ip_ready_substitutes(self): context = self._make_context(router_ip="10.0.0.5") result = _ProcessResult() @@ -3579,7 +3579,7 @@ async def test_replica_ip_ready_substitutes(self): assert out is not None assert context.job.job_spec.commands == ["echo 10.0.0.5"] - @freeze_time("2023-01-01 12:00:00Z") + @time_machine.travel(datetime(2023, 1, 1, 12, tzinfo=timezone.utc), tick=False) async def test_nodes_ref_in_service_terminates(self): context = self._make_context( router_ip="10.0.0.5", diff --git a/src/tests/_internal/server/background/pipeline_tasks/test_runs/test_termination.py b/src/tests/_internal/server/background/pipeline_tasks/test_runs/test_termination.py index 69016a915c..3a4eec04b1 100644 --- a/src/tests/_internal/server/background/pipeline_tasks/test_runs/test_termination.py +++ b/src/tests/_internal/server/background/pipeline_tasks/test_runs/test_termination.py @@ -4,7 +4,7 @@ from unittest.mock import AsyncMock, patch import pytest -from freezegun import freeze_time +import time_machine from sqlalchemy.ext.asyncio import AsyncSession from dstack._internal.core.models.configurations import TaskConfiguration @@ -187,7 +187,7 @@ async def test_finishes_non_scheduled_run_when_all_jobs_are_finished( assert run.lock_expires_at is None assert run.lock_owner is None - @freeze_time(datetime(2023, 1, 2, 3, 10, tzinfo=timezone.utc)) + @time_machine.travel(datetime(2023, 1, 2, 3, 10, tzinfo=timezone.utc), tick=False) async def test_reschedules_scheduled_run_and_clears_fleet( self, test_db, session: AsyncSession, worker: RunWorker ) -> None: diff --git a/src/tests/_internal/server/background/scheduled_tasks/test_events.py b/src/tests/_internal/server/background/scheduled_tasks/test_events.py index 91eb066f58..db8a6286ef 100644 --- a/src/tests/_internal/server/background/scheduled_tasks/test_events.py +++ b/src/tests/_internal/server/background/scheduled_tasks/test_events.py @@ -2,7 +2,7 @@ from unittest.mock import patch import pytest -from freezegun import freeze_time +import time_machine from sqlalchemy.ext.asyncio import AsyncSession from dstack._internal.server import settings @@ -16,7 +16,7 @@ async def test_deletes_old_events(test_db, session: AsyncSession) -> None: user = await create_user(session=session) for i in range(10): - with freeze_time(datetime(2026, 1, 1, i)): + with time_machine.travel(datetime(2026, 1, 1, i), tick=False): events.emit( session, message=f"Event {i}", @@ -30,7 +30,7 @@ async def test_deletes_old_events(test_db, session: AsyncSession) -> None: with ( patch.multiple(settings, SERVER_EVENTS_TTL_SECONDS=5 * 3600), - freeze_time(datetime(2026, 1, 1, 10)), + time_machine.travel(datetime(2026, 1, 1, 10), tick=False), ): await delete_events() diff --git a/src/tests/_internal/server/background/scheduled_tasks/test_metrics.py b/src/tests/_internal/server/background/scheduled_tasks/test_metrics.py index 1e3900a449..408f85ba61 100644 --- a/src/tests/_internal/server/background/scheduled_tasks/test_metrics.py +++ b/src/tests/_internal/server/background/scheduled_tasks/test_metrics.py @@ -2,7 +2,7 @@ from unittest.mock import patch import pytest -from freezegun import freeze_time +import time_machine from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession @@ -93,7 +93,7 @@ async def test_collects_metrics(self, test_db, session: AsyncSession): class TestDeleteMetrics: @pytest.mark.asyncio @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) - @freeze_time(datetime(2023, 1, 2, 3, 5, 20, tzinfo=timezone.utc)) + @time_machine.travel(datetime(2023, 1, 2, 3, 5, 20, tzinfo=timezone.utc), tick=False) async def test_deletes_old_metrics_running_job(self, test_db, session: AsyncSession): user = await create_user(session=session, global_role=GlobalRole.USER) project = await create_project(session=session, owner=user) @@ -141,7 +141,7 @@ async def test_deletes_old_metrics_running_job(self, test_db, session: AsyncSess @pytest.mark.asyncio @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) - @freeze_time(datetime(2023, 1, 2, 3, 5, 20, tzinfo=timezone.utc)) + @time_machine.travel(datetime(2023, 1, 2, 3, 5, 20, tzinfo=timezone.utc), tick=False) async def test_deletes_old_metrics_finished_job(self, test_db, session: AsyncSession): user = await create_user(session=session, global_role=GlobalRole.USER) project = await create_project(session=session, owner=user) diff --git a/src/tests/_internal/server/background/scheduled_tasks/test_probes.py b/src/tests/_internal/server/background/scheduled_tasks/test_probes.py index 22a9ba3383..f441046476 100644 --- a/src/tests/_internal/server/background/scheduled_tasks/test_probes.py +++ b/src/tests/_internal/server/background/scheduled_tasks/test_probes.py @@ -2,7 +2,7 @@ from unittest.mock import patch import pytest -from freezegun import freeze_time +import time_machine from sqlalchemy.ext.asyncio import AsyncSession from dstack._internal.core.models.configurations import ProbeConfig, ServiceConfiguration @@ -138,7 +138,7 @@ async def test_schedules_probe_execution(self, test_db, session: AsyncSession) - session, job, probe_num=1, due=datetime(2025, 1, 1, 0, 0, 0, tzinfo=timezone.utc) ) processing_time = datetime(2025, 1, 1, 0, 0, 1, tzinfo=timezone.utc) - with freeze_time(processing_time): + with time_machine.travel(processing_time, tick=False): with patch( "dstack._internal.server.background.scheduled_tasks.probes.PROBES_SCHEDULER" ) as scheduler_mock: diff --git a/src/tests/_internal/server/background/scheduled_tasks/test_prometheus_metrics.py b/src/tests/_internal/server/background/scheduled_tasks/test_prometheus_metrics.py index 0775723b4d..d0e641abf6 100644 --- a/src/tests/_internal/server/background/scheduled_tasks/test_prometheus_metrics.py +++ b/src/tests/_internal/server/background/scheduled_tasks/test_prometheus_metrics.py @@ -4,7 +4,7 @@ import pytest import pytest_asyncio -from freezegun import freeze_time +import time_machine from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession @@ -83,7 +83,7 @@ def shim_client_mock(self) -> Generator[Mock, None, None]: ) as ShimClientMock: yield ShimClientMock.return_value - @freeze_time(datetime(2023, 1, 2, 3, 5, 20, tzinfo=timezone.utc)) + @time_machine.travel(datetime(2023, 1, 2, 3, 5, 20, tzinfo=timezone.utc), tick=False) async def test_inserts_new_record( self, session: AsyncSession, job: JobModel, ssh_tunnel_mock: Mock, shim_client_mock: Mock ): @@ -100,7 +100,7 @@ async def test_inserts_new_record( assert metrics.text == "# prom response" assert metrics.collected_at == datetime(2023, 1, 2, 3, 5, 20, tzinfo=timezone.utc) - @freeze_time(datetime(2023, 1, 2, 3, 5, 20, tzinfo=timezone.utc)) + @time_machine.travel(datetime(2023, 1, 2, 3, 5, 20, tzinfo=timezone.utc), tick=False) async def test_updates_record( self, session: AsyncSession, job: JobModel, ssh_tunnel_mock: Mock, shim_client_mock: Mock ): @@ -125,7 +125,7 @@ async def test_updates_record( assert metrics.text == "# prom new response" assert metrics.collected_at == datetime(2023, 1, 2, 3, 5, 20, tzinfo=timezone.utc) - @freeze_time(datetime(2023, 1, 2, 3, 5, 20, tzinfo=timezone.utc)) + @time_machine.travel(datetime(2023, 1, 2, 3, 5, 20, tzinfo=timezone.utc), tick=False) async def test_skips_recently_updated( self, session: AsyncSession, job: JobModel, ssh_tunnel_mock: Mock, shim_client_mock: Mock ): @@ -150,7 +150,7 @@ async def test_skips_recently_updated( assert metrics.text == "# prom old response" assert metrics.collected_at == datetime(2023, 1, 2, 3, 5, 15, tzinfo=timezone.utc) - @freeze_time(datetime(2023, 1, 2, 3, 5, 20, tzinfo=timezone.utc)) + @time_machine.travel(datetime(2023, 1, 2, 3, 5, 20, tzinfo=timezone.utc), tick=False) @pytest.mark.dockerized(False) async def test_skips_non_dockerized_jobs( self, session: AsyncSession, job: JobModel, ssh_tunnel_mock: Mock, shim_client_mock: Mock @@ -170,7 +170,7 @@ async def test_skips_non_dockerized_jobs( @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) @pytest.mark.usefixtures("test_db", "image_config_mock") class TestDeletePrometheusMetrics: - @freeze_time(datetime(2023, 1, 2, 3, 5, 20, tzinfo=timezone.utc)) + @time_machine.travel(datetime(2023, 1, 2, 3, 5, 20, tzinfo=timezone.utc), tick=False) async def test_deletes_old_metrics(self, session: AsyncSession): user = await create_user(session=session, global_role=GlobalRole.USER) project = await create_project(session=session, owner=user) diff --git a/src/tests/_internal/server/conftest.py b/src/tests/_internal/server/conftest.py index 43ee83b736..9459c377c4 100644 --- a/src/tests/_internal/server/conftest.py +++ b/src/tests/_internal/server/conftest.py @@ -19,33 +19,6 @@ ) -def _warm_up_route_schemas() -> None: - """ - Build every route's pydantic schemas once, at import time. - - FastAPI builds a route's dependant lazily, on the first request that matches it. Several tests - make that first request inside `@freeze_time`, where `datetime.datetime` is freezegun's - `FakeDatetime`. pydantic v2 matches `datetime` by exact type and rejects the subclass, so a - route with an `Optional[datetime]` query parameter (`after`, `before`) fails to build with - `PydanticSchemaGenerationError` — and because `FakeDatetime` mimics `datetime`'s repr, the - error names `datetime.datetime` and reads like a production bug. - - Doing this eagerly, before any test freezes the clock, keeps the schemas built from the real - types. Nothing in production freezes time, so this is a test-environment fix only. - """ - # FastAPI 0.141 defers this to `_IncludedRouter.effective_candidates()`, reached from - # `matches()` on the first request. Duck-typed rather than importing the private class. - pending = list(app.routes) - while pending: - route = pending.pop() - build = getattr(route, "effective_candidates", None) - if callable(build): - pending.extend(build()) - - -_warm_up_route_schemas() - - @pytest.fixture def client(): transport = httpx.ASGITransport(app=app) diff --git a/src/tests/_internal/server/routers/test_events.py b/src/tests/_internal/server/routers/test_events.py index 0b07e70fdc..1c9149791c 100644 --- a/src/tests/_internal/server/routers/test_events.py +++ b/src/tests/_internal/server/routers/test_events.py @@ -4,7 +4,7 @@ import pytest import pytest_asyncio -from freezegun import freeze_time +import time_machine from httpx import AsyncClient from sqlalchemy import delete from sqlalchemy.ext.asyncio import AsyncSession @@ -49,14 +49,14 @@ async def test_response_format(self, session: AsyncSession, client: AsyncClient) ) event_ids = [uuid.uuid4() for _ in range(2)] with patch("uuid.uuid4", side_effect=event_ids): - with freeze_time(datetime(2026, 1, 1, 12, 0, 0)): + with time_machine.travel(datetime(2026, 1, 1, 12, 0, 0), tick=False): events.emit( session, "User added to project", actor=events.UserActor.from_user(user), targets=[events.Target.from_model(user), events.Target.from_model(project)], ) - with freeze_time(datetime(2026, 1, 1, 12, 0, 1)): + with time_machine.travel(datetime(2026, 1, 1, 12, 0, 1), tick=False): events.emit( session, "Project updated", @@ -1409,7 +1409,7 @@ async def test_pagination( for i in range(5): user = await create_user(session=session, name=f"user_{i}") users.append(user) - with freeze_time(datetime(2026, 1, 1, 12, 0, 0, i)): + with time_machine.travel(datetime(2026, 1, 1, 12, 0, 0, i), tick=False): events.emit( session, "User created", @@ -1480,14 +1480,14 @@ async def test_limits_events_regardless_number_of_targets( self, session: AsyncSession, client: AsyncClient ) -> None: users = [await create_user(session=session, name=f"user_{i}") for i in range(3)] - with freeze_time(datetime(2026, 1, 1, 12, 0, 0, 0)): + with time_machine.travel(datetime(2026, 1, 1, 12, 0, 0, 0), tick=False): events.emit( session, "Users batch created", actor=events.SystemActor(), targets=[events.Target.from_model(users[0]), events.Target.from_model(users[1])], ) - with freeze_time(datetime(2026, 1, 1, 12, 0, 0, 1)): + with time_machine.travel(datetime(2026, 1, 1, 12, 0, 0, 1), tick=False): events.emit( session, "User created", diff --git a/src/tests/_internal/server/routers/test_fleets.py b/src/tests/_internal/server/routers/test_fleets.py index eb817d9362..4182126199 100644 --- a/src/tests/_internal/server/routers/test_fleets.py +++ b/src/tests/_internal/server/routers/test_fleets.py @@ -5,7 +5,7 @@ from uuid import uuid4 import pytest -from freezegun import freeze_time +import time_machine from httpx import AsyncClient from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession @@ -937,7 +937,7 @@ async def test_returns_40x_if_not_authenticated(self, client: AsyncClient): @pytest.mark.asyncio @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) - @freeze_time(datetime(2023, 1, 2, 3, 4, tzinfo=timezone.utc)) + @time_machine.travel(datetime(2023, 1, 2, 3, 4, tzinfo=timezone.utc), tick=False) async def test_creates_fleet(self, test_db, session: AsyncSession, client: AsyncClient): user = await create_user(session, global_role=GlobalRole.USER) project = await create_project(session) @@ -1044,7 +1044,7 @@ async def test_creates_fleet(self, test_db, session: AsyncSession, client: Async @pytest.mark.asyncio @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) - @freeze_time(datetime(2023, 1, 2, 3, 4, tzinfo=timezone.utc)) + @time_machine.travel(datetime(2023, 1, 2, 3, 4, tzinfo=timezone.utc), tick=False) async def test_creates_ssh_fleet(self, test_db, session: AsyncSession, client: AsyncClient): user = await create_user(session, global_role=GlobalRole.USER) project = await create_project(session) @@ -1228,7 +1228,7 @@ async def test_creates_ssh_fleet_with_blocks( @pytest.mark.asyncio @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) - @freeze_time(datetime(2023, 1, 2, 3, 4, tzinfo=timezone.utc), real_asyncio=True) + @time_machine.travel(datetime(2023, 1, 2, 3, 4, tzinfo=timezone.utc), tick=False) async def test_updates_ssh_fleet(self, test_db, session: AsyncSession, client: AsyncClient): user = await create_user(session, global_role=GlobalRole.USER) project = await create_project(session) @@ -1540,7 +1540,7 @@ async def test_updates_cloud_fleet_nodes_target_without_changing_instance_count( @pytest.mark.asyncio @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) - @freeze_time(datetime(2023, 1, 2, 3, 4, tzinfo=timezone.utc)) + @time_machine.travel(datetime(2023, 1, 2, 3, 4, tzinfo=timezone.utc), tick=False) async def test_errors_if_ssh_key_is_bad( self, test_db, session: AsyncSession, client: AsyncClient ): diff --git a/src/tests/_internal/server/routers/test_projects.py b/src/tests/_internal/server/routers/test_projects.py index ff25357515..7e7be9d0f6 100644 --- a/src/tests/_internal/server/routers/test_projects.py +++ b/src/tests/_internal/server/routers/test_projects.py @@ -3,7 +3,7 @@ from uuid import UUID import pytest -from freezegun import freeze_time +import time_machine from httpx import AsyncClient from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession @@ -928,7 +928,7 @@ async def test_returns_40x_if_not_authenticated(self, client: AsyncClient): @pytest.mark.asyncio @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) - @freeze_time(datetime(2023, 1, 2, 3, 4, tzinfo=timezone.utc)) + @time_machine.travel(datetime(2023, 1, 2, 3, 4, tzinfo=timezone.utc), tick=False) async def test_creates_project(self, test_db, session: AsyncSession, client: AsyncClient): user = await create_user(session=session) project_id = UUID("1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e") @@ -1070,7 +1070,7 @@ async def test_forbids_if_no_permission_to_create_projects( @pytest.mark.asyncio @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) - @freeze_time(datetime(2023, 1, 2, 3, 4, tzinfo=timezone.utc)) + @time_machine.travel(datetime(2023, 1, 2, 3, 4, tzinfo=timezone.utc), tick=False) async def test_creates_public_project( self, test_db, session: AsyncSession, client: AsyncClient ): @@ -1099,7 +1099,7 @@ async def test_creates_public_project( @pytest.mark.asyncio @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) - @freeze_time(datetime(2023, 1, 2, 3, 4, tzinfo=timezone.utc)) + @time_machine.travel(datetime(2023, 1, 2, 3, 4, tzinfo=timezone.utc), tick=False) async def test_creates_private_project_by_default( self, test_db, session: AsyncSession, client: AsyncClient ): @@ -1128,7 +1128,7 @@ async def test_creates_private_project_by_default( @pytest.mark.asyncio @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) - @freeze_time(datetime(2023, 1, 2, 3, 4, tzinfo=timezone.utc)) + @time_machine.travel(datetime(2023, 1, 2, 3, 4, tzinfo=timezone.utc), tick=False) async def test_creates_private_project_explicitly( self, test_db, session: AsyncSession, client: AsyncClient ): diff --git a/src/tests/_internal/server/routers/test_prometheus.py b/src/tests/_internal/server/routers/test_prometheus.py index 8b8910273e..26da89a90e 100644 --- a/src/tests/_internal/server/routers/test_prometheus.py +++ b/src/tests/_internal/server/routers/test_prometheus.py @@ -4,7 +4,6 @@ from unittest.mock import patch import pytest -from freezegun import freeze_time from httpx import AsyncClient from sqlalchemy.ext.asyncio import AsyncSession @@ -67,7 +66,7 @@ def enable_metrics(monkeypatch: pytest.MonkeyPatch): FAKE_NOW = datetime(2023, 1, 2, 3, 4, tzinfo=timezone.utc) -@freeze_time(FAKE_NOW) +@pytest.mark.time_machine(FAKE_NOW, tick=False) @pytest.mark.asyncio @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) @pytest.mark.usefixtures("image_config_mock", "test_db", "enable_metrics") diff --git a/src/tests/_internal/server/routers/test_public_keys.py b/src/tests/_internal/server/routers/test_public_keys.py index e4dee3895b..c344d9f13f 100644 --- a/src/tests/_internal/server/routers/test_public_keys.py +++ b/src/tests/_internal/server/routers/test_public_keys.py @@ -3,7 +3,7 @@ from unittest.mock import AsyncMock import pytest -from freezegun import freeze_time +import time_machine from httpx import AsyncClient from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession @@ -118,7 +118,7 @@ async def test_returns_40x_if_not_authenticated(self, client: AsyncClient): @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) @pytest.mark.usefixtures("test_db") - @freeze_time(datetime(2023, 1, 2, 3, 4, tzinfo=timezone.utc)) + @time_machine.travel(datetime(2023, 1, 2, 3, 4, tzinfo=timezone.utc), tick=False) async def test_adds_valid_public_key( self, session: AsyncSession, diff --git a/src/tests/_internal/server/routers/test_runs.py b/src/tests/_internal/server/routers/test_runs.py index 5ba350efb8..e16c81cea7 100644 --- a/src/tests/_internal/server/routers/test_runs.py +++ b/src/tests/_internal/server/routers/test_runs.py @@ -6,7 +6,7 @@ from uuid import UUID import pytest -from freezegun import freeze_time +import time_machine from httpx import AsyncClient from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession @@ -3571,7 +3571,7 @@ async def test_creates_pending_run_if_run_is_scheduled( repo_id=repo.name, ) run_spec.configuration.schedule = Schedule(cron=["5 * * * *", "10 * * * *"]) - with freeze_time(datetime(2023, 1, 2, 3, 9, tzinfo=timezone.utc)): + with time_machine.travel(datetime(2023, 1, 2, 3, 9, tzinfo=timezone.utc), tick=False): response = await client.post( f"/api/project/{project.name}/runs/apply", headers=get_auth_headers(user.token), diff --git a/src/tests/_internal/server/routers/test_users.py b/src/tests/_internal/server/routers/test_users.py index ac7ba24207..b6c5c24485 100644 --- a/src/tests/_internal/server/routers/test_users.py +++ b/src/tests/_internal/server/routers/test_users.py @@ -3,7 +3,7 @@ from uuid import UUID import pytest -from freezegun import freeze_time +import time_machine from httpx import AsyncClient from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession @@ -413,7 +413,7 @@ async def test_returns_40x_if_not_authenticated(self, client: AsyncClient): @pytest.mark.asyncio @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) - @freeze_time(datetime(2023, 1, 2, 3, 4, tzinfo=timezone.utc)) + @time_machine.travel(datetime(2023, 1, 2, 3, 4, tzinfo=timezone.utc), tick=False) async def test_creates_user(self, test_db, session: AsyncSession, client: AsyncClient): user = await create_user(name="admin", session=session) with patch("uuid.uuid4") as uuid_mock: @@ -448,7 +448,7 @@ async def test_creates_user(self, test_db, session: AsyncSession, client: AsyncC @pytest.mark.asyncio @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) - @freeze_time(datetime(2023, 1, 2, 3, 4, tzinfo=timezone.utc)) + @time_machine.travel(datetime(2023, 1, 2, 3, 4, tzinfo=timezone.utc), tick=False) async def test_return_400_if_username_taken( self, test_db, session: AsyncSession, client: AsyncClient ): @@ -501,7 +501,7 @@ async def test_return_400_if_username_taken( @pytest.mark.asyncio @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) - @freeze_time(datetime(2023, 1, 2, 3, 4, tzinfo=timezone.utc)) + @time_machine.travel(datetime(2023, 1, 2, 3, 4, tzinfo=timezone.utc), tick=False) async def test_returns_400_if_username_invalid( self, test_db, diff --git a/src/tests/_internal/server/routers/test_volumes.py b/src/tests/_internal/server/routers/test_volumes.py index e007f928c0..30cab24233 100644 --- a/src/tests/_internal/server/routers/test_volumes.py +++ b/src/tests/_internal/server/routers/test_volumes.py @@ -4,7 +4,7 @@ from uuid import UUID import pytest -from freezegun import freeze_time +import time_machine from httpx import AsyncClient from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession @@ -308,7 +308,7 @@ async def test_returns_40x_if_not_authenticated(self, client: AsyncClient): @pytest.mark.asyncio @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) - @freeze_time(datetime(2023, 1, 2, 3, 4, tzinfo=timezone.utc)) + @time_machine.travel(datetime(2023, 1, 2, 3, 4, tzinfo=timezone.utc), tick=False) async def test_creates_volume(self, test_db, session: AsyncSession, client: AsyncClient): user = await create_user(session, global_role=GlobalRole.USER) project = await create_project(session) diff --git a/src/tests/_internal/server/services/test_logs.py b/src/tests/_internal/server/services/test_logs.py index 8a8cf3394f..5154f903f1 100644 --- a/src/tests/_internal/server/services/test_logs.py +++ b/src/tests/_internal/server/services/test_logs.py @@ -8,7 +8,7 @@ import botocore.exceptions import pytest import pytest_asyncio -from freezegun import freeze_time +import time_machine from sqlalchemy.ext.asyncio import AsyncSession from dstack._internal.core.errors import ServerClientError @@ -799,7 +799,7 @@ class TestCloudWatchLogStorage: FAKE_NOW = datetime(2023, 10, 6, 10, 1, 54, tzinfo=timezone.utc) @pytest_asyncio.fixture - @freeze_time(FAKE_NOW) + @time_machine.travel(FAKE_NOW, tick=False) async def project(self, test_db, session: AsyncSession) -> ProjectModel: project = await create_project(session=session, name="test-proj") return project @@ -1163,7 +1163,7 @@ async def test_poll_logs_exception_other( log_storage.poll_logs(project, poll_logs_request) @pytest.mark.asyncio - @freeze_time(FAKE_NOW) + @time_machine.travel(FAKE_NOW, tick=False) async def test_write_logs( self, project: ProjectModel, @@ -1215,7 +1215,7 @@ async def test_write_logs( mock_client.put_log_events.assert_has_calls(expected_put_log_events_calls, any_order=True) @pytest.mark.asyncio - @freeze_time(FAKE_NOW) + @time_machine.travel(FAKE_NOW, tick=False) async def test_write_logs_resource_not_found( self, project: ProjectModel, @@ -1250,7 +1250,7 @@ async def test_write_logs_resource_not_found( assert mock_client.put_log_events.call_count == 2 @pytest.mark.asyncio - @freeze_time(FAKE_NOW) + @time_machine.travel(FAKE_NOW, tick=False) async def test_write_logs_other_exception( self, project: ProjectModel, @@ -1271,7 +1271,7 @@ async def test_write_logs_other_exception( ) @pytest.mark.asyncio - @freeze_time(FAKE_NOW) + @time_machine.travel(FAKE_NOW, tick=False) async def test_write_logs_not_in_chronological_order( self, caplog: pytest.LogCaptureFixture, @@ -1309,7 +1309,7 @@ async def test_write_logs_not_in_chronological_order( assert "events are not in chronological order" in caplog.text @pytest.mark.asyncio - @freeze_time(FAKE_NOW) + @time_machine.travel(FAKE_NOW, tick=False) async def test_write_logs_past_and_future_events( self, caplog: pytest.LogCaptureFixture, @@ -1375,7 +1375,7 @@ def _delta_ms(**kwargs: int) -> int: ], ) @pytest.mark.asyncio - @freeze_time(FAKE_NOW) + @time_machine.travel(FAKE_NOW, tick=False) async def test_write_logs_batching_by_size( self, monkeypatch: pytest.MonkeyPatch, @@ -1422,7 +1422,7 @@ async def test_write_logs_batching_by_size( ], ) @pytest.mark.asyncio - @freeze_time(FAKE_NOW) + @time_machine.travel(FAKE_NOW, tick=False) async def test_write_logs_batching_by_count( self, monkeypatch: pytest.MonkeyPatch, @@ -1454,7 +1454,7 @@ async def test_write_logs_batching_by_count( assert actual == expected @pytest.mark.asyncio - @freeze_time(FAKE_NOW) + @time_machine.travel(FAKE_NOW, tick=False) async def test_write_logs_batching_by_timestamp( self, project: ProjectModel, diff --git a/src/tests/_internal/server/services/test_volumes.py b/src/tests/_internal/server/services/test_volumes.py index b05455ed31..e77e2591db 100644 --- a/src/tests/_internal/server/services/test_volumes.py +++ b/src/tests/_internal/server/services/test_volumes.py @@ -1,7 +1,7 @@ from datetime import datetime, timezone import pytest -from freezegun import freeze_time +import time_machine from dstack._internal.core.errors import ServerClientError from dstack._internal.core.models.backends.base import BackendType @@ -115,7 +115,7 @@ def test_returns_0_when_no_price(self): ) assert _get_volume_cost(volume) == 0.0 - @freeze_time(datetime(2025, 1, 31, 0, 0, tzinfo=timezone.utc)) + @time_machine.travel(datetime(2025, 1, 31, 0, 0, tzinfo=timezone.utc), tick=False) def test_calculates_active_volume_cost(self): volume = get_volume( status=VolumeStatus.ACTIVE, @@ -125,7 +125,7 @@ def test_calculates_active_volume_cost(self): ) assert _get_volume_cost(volume) == pytest.approx(30.0) - @freeze_time(datetime(2025, 1, 31, 0, 0, tzinfo=timezone.utc)) + @time_machine.travel(datetime(2025, 1, 31, 0, 0, tzinfo=timezone.utc), tick=False) def test_calculates_finished_volume_cost(self): volume = get_volume( provisioning_data=get_volume_provisioning_data(price=30), @@ -136,7 +136,7 @@ def test_calculates_finished_volume_cost(self): # Cost should be for 15 days out of a 30-day pricing period assert _get_volume_cost(volume) == pytest.approx(15.0) - @freeze_time(datetime(2025, 1, 1, 0, 0, tzinfo=timezone.utc)) + @time_machine.travel(datetime(2025, 1, 1, 0, 0, tzinfo=timezone.utc), tick=False) def test_calculates_zero_cost_for_zero_duration_active(self): volume = get_volume( status=VolumeStatus.ACTIVE, diff --git a/src/tests/_internal/utils/test_common.py b/src/tests/_internal/utils/test_common.py index 6f26151bdc..91f8bf0768 100644 --- a/src/tests/_internal/utils/test_common.py +++ b/src/tests/_internal/utils/test_common.py @@ -3,7 +3,6 @@ from typing import Any, Iterable import pytest -from freezegun import freeze_time from dstack._internal.utils.common import ( batched, @@ -31,7 +30,7 @@ def test_local_time(dt: datetime, result: str) -> None: assert local_time(dt) == result -@freeze_time(datetime(2023, 10, 4, 12, 0, tzinfo=timezone.utc)) +@pytest.mark.time_machine(datetime(2023, 10, 4, 12, 0, tzinfo=timezone.utc), tick=False) class TestPrettyDate: def test_now(self): now = datetime.now(tz=timezone.utc)