From 1e0b397720899ca1663038de08e3c6bd3d62edc2 Mon Sep 17 00:00:00 2001
From: William Chu <william.chu@uptickhq.com>
Date: Wed, 14 Feb 2024 17:39:09 +1100
Subject: [PATCH] fix(timestamps): generate timestamps correctly and remove the
 deprecated pytz

---
 backend/bot/scheduler/scheduler.py |  4 +--
 backend/bot/shared/tools.py        | 17 ++++++++----
 backend/requirements.txt           |  1 -
 backend/tests/test_tools.py        | 44 ++++++++++++++++++++++++++++++
 4 files changed, 57 insertions(+), 9 deletions(-)
 create mode 100644 backend/tests/test_tools.py

diff --git a/backend/bot/scheduler/scheduler.py b/backend/bot/scheduler/scheduler.py
index 4d8133f0..3fcbba1c 100644
--- a/backend/bot/scheduler/scheduler.py
+++ b/backend/bot/scheduler/scheduler.py
@@ -1,3 +1,4 @@
+import zoneinfo
 import config
 import datetime
 import slack_sdk
@@ -14,7 +15,6 @@
     store_slack_user_list_db,
 )
 from iblog import logger
-from pytz import timezone
 from typing import List
 
 
@@ -26,7 +26,7 @@ class TaskScheduler:
     def __init__(self):
         self.scheduler = BackgroundScheduler(
             jobstores=jobstores,
-            timezone=timezone(application_timezone),
+            timezone=zoneinfo.ZoneInfo(application_timezone),
         )
 
     def list_jobs(self) -> List[Job]:
diff --git a/backend/bot/shared/tools.py b/backend/bot/shared/tools.py
index 2f3dad9b..027fac33 100644
--- a/backend/bot/shared/tools.py
+++ b/backend/bot/shared/tools.py
@@ -1,4 +1,5 @@
 import config
+import zoneinfo
 import ipaddress
 import itertools
 import random
@@ -6,7 +7,6 @@
 
 from datetime import datetime
 from iblog import logger
-from pytz import timezone
 from typing import Any, List
 
 
@@ -19,18 +19,23 @@
 application_timezone = config.active.options.get("timezone")
 
 
-def fetch_timestamp(short: bool = False):
+def fetch_timestamp(short: bool = False, timezone: str | None = None) -> str:
     """Return a localized, formatted timestamp using datetime.now()"""
-    now = datetime.now()
-    localized = timezone(application_timezone).localize(now)
+    localized = datetime.now(
+        zoneinfo.ZoneInfo(timezone or application_timezone)
+    )
     if short:
         return localized.strftime(timestamp_fmt_short)
     return localized.strftime(timestamp_fmt)
 
 
-def fetch_timestamp_from_time_obj(t: datetime):
+def fetch_timestamp_from_time_obj(
+    t: datetime, timezone: str | None = None
+) -> str:
     """Return a localized, formatted timestamp using datetime.datetime class"""
-    return timezone(application_timezone).localize(t).strftime(timestamp_fmt)
+    return t.astimezone(
+        zoneinfo.ZoneInfo(timezone or application_timezone)
+    ).strftime(timestamp_fmt)
 
 
 def find_index_in_list(lst: List, key: Any, value: Any):
diff --git a/backend/requirements.txt b/backend/requirements.txt
index c57fdc2b..b7a24e1a 100644
--- a/backend/requirements.txt
+++ b/backend/requirements.txt
@@ -22,7 +22,6 @@ pytest-mock==3.14.0
 pytest-sqlalchemy-mock==0.1.5
 python-dotenv==1.0.1
 python-json-logger==2.0.7
-pytz==2024.1
 PyYAML==6.0.1
 requests==2.31.0
 requests-mock==1.11.0
diff --git a/backend/tests/test_tools.py b/backend/tests/test_tools.py
new file mode 100644
index 00000000..fe93f1f6
--- /dev/null
+++ b/backend/tests/test_tools.py
@@ -0,0 +1,44 @@
+import datetime
+
+from bot.shared.tools import (
+    fetch_timestamp,
+    fetch_timestamp_from_time_obj,
+    timestamp_fmt,
+    timestamp_fmt_short,
+)
+
+
+class TestTimeHelpers:
+    def test_fetch_timestamp_generates_the_correct_timestamp(self):
+        utc_timestamp = fetch_timestamp(timezone="UTC")
+        melbourne_timestamp = fetch_timestamp(timezone="Australia/Melbourne")
+
+        parsed_melbourne = datetime.datetime.strptime(
+            utc_timestamp[:13], "%Y-%m-%dT%H"
+        )
+        parsed_utc = datetime.datetime.strptime(
+            melbourne_timestamp[:13], "%Y-%m-%dT%H"
+        )
+
+        assert (
+            parsed_melbourne.hour != parsed_utc.hour
+        ), "The timestamps should have different hours when using differnt timezones"
+
+    def test_fetch_timestamp_short(self):
+        short_timestamp = fetch_timestamp(short=True, timezone="UTC")
+
+        assert datetime.datetime.strptime(
+            short_timestamp, timestamp_fmt_short
+        ), "Should be able to reparse the short format"
+
+    def test_fetch_timestamp_from_time_obj(self):
+        now = datetime.datetime.now(datetime.UTC)
+
+        now_as_est = fetch_timestamp_from_time_obj(
+            now, timezone="Australia/Melbourne"
+        )
+
+        now_as_melbourne_datetime = datetime.datetime.strptime(
+            now_as_est, timestamp_fmt
+        )
+        assert now_as_melbourne_datetime.hour != now.hour