From dc583a8155c0e2f20017cf996f6aeded94baedb1 Mon Sep 17 00:00:00 2001
From: Matthew Richards <matthew.richards@stfc.ac.uk>
Date: Fri, 14 May 2021 14:48:09 +0000
Subject: [PATCH]  #225: Fix timezone-related issues found by merging master

- These issues were a result of writing new code that involved timezones before master got merged in, which had the original fixes for timezones on DB backend
---
 datagateway_api/common/helpers.py | 5 +++++
 test/icat/test_query.py           | 1 -
 test/test_map_distinct_attrs.py   | 5 +++--
 3 files changed, 8 insertions(+), 3 deletions(-)

diff --git a/datagateway_api/common/helpers.py b/datagateway_api/common/helpers.py
index 452be102..88d8f178 100644
--- a/datagateway_api/common/helpers.py
+++ b/datagateway_api/common/helpers.py
@@ -3,6 +3,7 @@
 import json
 import logging
 
+from dateutil.tz.tz import tzlocal
 from flask import request
 from flask_restful import reqparse
 from sqlalchemy.exc import IntegrityError
@@ -157,6 +158,10 @@ def map_distinct_attributes_to_results(distinct_attributes, query_result):
         split_attr_name = attr_name.split(".")
 
         if isinstance(data, datetime):
+            # Workaround for when this function is used on DB backend, where usually
+            # `_make_serializable()` would fix tzinfo
+            if data.tzinfo is None:
+                data = data.replace(tzinfo=tzlocal())
             data = DateHandler.datetime_object_to_str(data)
 
         # Attribute name is from the 'origin' entity (i.e. not a related entity)
diff --git a/test/icat/test_query.py b/test/icat/test_query.py
index a4dd9fac..47df71f6 100644
--- a/test/icat/test_query.py
+++ b/test/icat/test_query.py
@@ -1,6 +1,5 @@
 from datetime import datetime
 
-from dateutil.tz import tzlocal
 from icat.entity import Entity
 import pytest
 
diff --git a/test/test_map_distinct_attrs.py b/test/test_map_distinct_attrs.py
index 626345ab..615bc3f5 100644
--- a/test/test_map_distinct_attrs.py
+++ b/test/test_map_distinct_attrs.py
@@ -1,5 +1,6 @@
-from datetime import datetime, timezone
+from datetime import datetime
 
+from dateutil.tz import tzlocal
 import pytest
 
 from datagateway_api.common.helpers import map_distinct_attributes_to_results
@@ -25,7 +26,7 @@ class TestMapDistinctAttrs:
                         hour=1,
                         minute=1,
                         second=1,
-                        tzinfo=timezone.utc,
+                        tzinfo=tzlocal(),
                     ),
                 ),
                 {"startDate": "2020-01-04 01:01:01+00:00"},