Fix flaky timeline writer test (#416)

Author: ndodda-amazon

tests/core/test_timeline_writer.py

@@ -268,31 +268,30 @@ def test_rotation_policy(rotation_profiler_config_parser, policy, out_dir):
 
 
 @pytest.mark.parametrize("timezone", ["Europe/Dublin", "Australia/Melbourne", "US/Eastern"])
-def test_utc_timestamp(simple_profiler_config_parser, timezone, out_dir):
+def test_utc_timestamp(monkeypatch, simple_profiler_config_parser, timezone, out_dir):
     """
     This test is meant to set to create files/events in different timezones and check if timeline writer stores
     them in UTC.
     """
+    monkeypatch.setenv("TZ", timezone)
     assert simple_profiler_config_parser.profiling_enabled
 
     time.tzset()
-    event_time_in_timezone = time.mktime(time.localtime())
-    time_in_utc = event_time_in_utc = calendar.timegm(time.gmtime())
 
     timeline_writer = TimelineFileWriter(profiler_config_parser=simple_profiler_config_parser)
     assert timeline_writer
 
     event_times_in_utc = []
     for i in range(1, 3):
+        event_time_in_timezone = time.mktime(time.localtime())
+        event_time_in_utc = time_in_utc = calendar.timegm(time.gmtime())
         event_times_in_utc.append(event_time_in_utc)
         timeline_writer.write_trace_events(
             training_phase=f"TimestampTest",
             op_name="event_in_" + timezone + str(i),
             timestamp=event_time_in_timezone,
             duration=20,
         )
-        event_time_in_timezone = time.mktime(time.localtime())
-        event_time_in_utc = calendar.timegm(time.gmtime())
 
     timeline_writer.flush()
     timeline_writer.close()
@@ -305,7 +304,7 @@ def test_utc_timestamp(simple_profiler_config_parser, timezone, out_dir):
     path = file_path.name.split(DEFAULT_PREFIX)
     file_timestamp = int(path[0].split("_")[0])
 
-    # file timestamp uses end of event
+    # file timestamp uses end of last event
     assert (time_in_utc + 20) * CONVERT_TO_MICROSECS == file_timestamp
 
     start_time_since_epoch = 0