Specify UTC timezone argument in datetime.now() and fix unit tests (#…

…6279)

While running unit tests locally using `make test`, I saw these
failures:
```
=========================================================================================================== short test summary info ===========================================================================================================
FAILED tests/clickhouse/optimize/test_optimize.py::test_optimize_partitions_raises_exception_with_cutoff_time - Failed: DID NOT RAISE <class 'snuba.clickhouse.optimize.optimize_scheduler.OptimizedSchedulerTimeout'>
FAILED tests/clickhouse/optimize/test_optimize_scheduler.py::test_get_next_schedule[non parallel] - assert OptimizationSchedule(partitions_groups=[["(90,'2022-06-27')",\n                                         "(90,'2022-06-20')",\n   ...
FAILED tests/clickhouse/optimize/test_optimize_scheduler.py::test_get_next_schedule[parallel before final cutoff] - assert OptimizationSchedule(partitions_groups=[["(90,'2022-03-28')",\n                                         "(90,'202...
FAILED tests/clickhouse/optimize/test_optimize_scheduler.py::test_get_next_schedule_raises_exception - Failed: DID NOT RAISE <class 'snuba.clickhouse.optimize.optimize_scheduler.OptimizedSchedulerTimeout'>
FAILED tests/subscriptions/test_scheduler_consumer.py::test_tick_time_shift - AssertionError: assert Tick(partition=0, offsets=Interval(lower=0, upper=1), timestamps=Interval(lower=86400.0, upper=172800.0)) == Tick(partition=0, offsets=...
========================================================================== 5 failed, 2601 passed, 4 skipped, 2 deselected, 1 xfailed, 6 xpassed in 430.69s (0:07:10) ==========================================================================
```

The reason for failure is that some of these tests use [time_machine's
context
manager](https://github.com/adamchainz/time-machine?tab=readme-ov-file#context-manager),
which expects the time in UTC. When these tests are run locally in PST,
there is a difference of 7 hours that never triggers the cutoff time
expiry.
```
>>> from datetime import datetime, timezone
>>> datetime.now()
datetime.datetime(2024, 9, 9, 11, 51, 15, 279890)
>>> datetime.now(timezone.utc)
datetime.datetime(2024, 9, 9, 18, 51, 26, 40580, tzinfo=datetime.timezone.utc)
```

The fix is to use `datetime`'s `timezone.utc` instead of
`datetime.now()`. `Datetime` also used to have `datetime.utcnow()` but
it has been deprecated.

Note: I also fixed some non-test code, so I would appreciate reviews
from experts in those areas.

snuba/clickhouse/optimize/optimize.py

@@ -5,7 +5,7 @@
 import time
 from collections import deque
 from concurrent.futures import Future, ThreadPoolExecutor
-from datetime import datetime, timedelta
+from datetime import datetime, timedelta, timezone
 from typing import Any, Mapping, Optional, Sequence
 
 import structlog
@@ -347,7 +347,7 @@ def optimize_partitions(
  """
 
  for partition in partitions:
- if cutoff_time is not None and datetime.now() > cutoff_time:
+ if cutoff_time is not None and datetime.now(timezone.utc) > cutoff_time:
  logger.info(
  f"Optimize job is running past provided cutoff time"
  f" {cutoff_time}. Cancelling.",

snuba/clickhouse/optimize/optimize_scheduler.py

@@ -1,6 +1,6 @@
 import re
 from dataclasses import dataclass
-from datetime import datetime, timedelta
+from datetime import datetime, timedelta, timezone
 from typing import MutableSequence, Sequence
 
 from snuba import settings
@@ -42,9 +42,9 @@ class OptimizeScheduler:
 
  def __init__(self, default_parallel_threads: int) -> None:
  self.__default_parallel_threads = default_parallel_threads
- self.__last_midnight = (datetime.now() + timedelta(minutes=10)).replace(
- hour=0, minute=0, second=0, microsecond=0
- )
+ self.__last_midnight = (
+ datetime.now(timezone.utc) + timedelta(minutes=10)
+ ).replace(hour=0, minute=0, second=0, microsecond=0)
  self.__parallel_start_time = self.__last_midnight + timedelta(
  hours=settings.PARALLEL_OPTIMIZE_JOB_START_TIME
  )
@@ -95,7 +95,7 @@ def get_next_schedule(self, partitions: Sequence[str]) -> OptimizationSchedule:
  reached.
  """
  num_threads = get_num_threads(self.__default_parallel_threads)
- current_time = datetime.now()
+ current_time = datetime.now(timezone.utc)
  if current_time >= self.__full_job_end_time:
  raise OptimizedSchedulerTimeout(
  f"Optimize job cutoff time exceeded "

snuba/settings/settings_test.py

@@ -40,7 +40,7 @@
 ENFORCE_RETENTION = True
 
 # Ignore optimize job cut off time for tests
-OPTIMIZE_JOB_CUTOFF_TIME = 24
+OPTIMIZE_JOB_CUTOFF_TIME = 23
 
 OPTIMIZE_PARALLEL_MAX_JITTER_MINUTES = 0
 

tests/clickhouse/optimize/test_optimize_scheduler.py

@@ -1,4 +1,4 @@
-from datetime import datetime, timedelta
+from datetime import datetime, timedelta, timezone
 from typing import Sequence
 
 import pytest
@@ -132,7 +132,7 @@ def test_subdivide_partitions(
  )
 
 
-last_midnight = (datetime.now() + timedelta(minutes=10)).replace(
+last_midnight = (datetime.now(timezone.utc) + timedelta(minutes=10)).replace(
  hour=0, minute=0, second=0, microsecond=0
 )
 

tests/clickhouse/optimize/test_optimize_tracker.py

@@ -1,6 +1,6 @@
 import time
 import uuid
-from datetime import datetime, timedelta
+from datetime import datetime, timedelta, timezone
 from typing import Optional, Set
 from unittest.mock import call, patch
 
@@ -136,11 +136,11 @@ def write_error_message(writable_storage: WritableTableStorage, time: int) -> No
  for week in range(0, 4):
  write_error_message(
  writable_storage=storage,
- time=int((datetime.now() - timedelta(weeks=week)).timestamp()),
+ time=int((datetime.now(timezone.utc) - timedelta(weeks=week)).timestamp()),
  )
  write_error_message(
  writable_storage=storage,
- time=int((datetime.now() - timedelta(weeks=week)).timestamp()),
+ time=int((datetime.now(timezone.utc) - timedelta(weeks=week)).timestamp()),
  )
 
  partitions = optimize.get_partitions_to_optimize(
@@ -227,11 +227,11 @@ def write_error_message(writable_storage: WritableTableStorage, time: int) -> No
  for week in range(0, 4):
  write_error_message(
  writable_storage=storage,
- time=int((datetime.now() - timedelta(weeks=week)).timestamp()),
+ time=int((datetime.now(timezone.utc) - timedelta(weeks=week)).timestamp()),
  )
  write_error_message(
  writable_storage=storage,
- time=int((datetime.now() - timedelta(weeks=week)).timestamp()),
+ time=int((datetime.now(timezone.utc) - timedelta(weeks=week)).timestamp()),
  )
 
  partitions = optimize.get_partitions_to_optimize(

tests/subscriptions/test_scheduler_consumer.py

@@ -2,7 +2,7 @@
 import logging
 import time
 import uuid
-from datetime import datetime, timedelta
+from datetime import datetime, timedelta, timezone
 from typing import Any, Mapping, Optional
 from unittest import mock
 
@@ -148,7 +148,10 @@ def test_tick_time_shift() -> None:
  assert tick.time_shift(timedelta(hours=24).total_seconds()) == Tick(
  partition,
  offsets,
- Interval(datetime(1970, 1, 2).timestamp(), datetime(1970, 1, 3).timestamp()),
+ Interval(
+ datetime(1970, 1, 2, tzinfo=timezone.utc).timestamp(),
+ datetime(1970, 1, 3, tzinfo=timezone.utc).timestamp(),
+ ),
  )
 
 
@@ -437,7 +440,7 @@ def test_invalid_commit_log_message(caplog: Any) -> None:
  "orig_message_ts",
  )
 
- now = datetime.now()
+ now = datetime.now(timezone.utc)
 
  def _assignment_callback(offsets: Mapping[Partition, int]) -> None:
  assert inner_consumer.tell() == {partition: 0}