Use the same source of time in flaky test (#3009)

* Use the same source of time in flaky test Signed-off-by: Eduardo Apolinario <[email protected]> * Use proto Timestamp to get current time Signed-off-by: Eduardo Apolinario <[email protected]> --------- Signed-off-by: Eduardo Apolinario <[email protected]> Co-authored-by: Eduardo Apolinario <[email protected]>
flyteorg · Dec 17, 2024 · 34af2e2 · 34af2e2
1 parent e5c2f41
commit 34af2e2
Showing 1 changed file with 9 additions and 6 deletions.
diff --git a/tests/flytekit/unit/bin/test_python_entrypoint.py b/tests/flytekit/unit/bin/test_python_entrypoint.py
@@ -526,17 +526,20 @@ def test_get_container_error_timestamp(monkeypatch) -> None:
 
     assert get_container_error_timestamp(FlyteException("foo", timestamp=10.5)) == Timestamp(seconds=10, nanos=500000000)
 
-    current_dtime = datetime.now()
+    current_timestamp = Timestamp()
+    current_timestamp.GetCurrentTime()
     error_timestamp = get_container_error_timestamp(RuntimeError("foo"))
-    assert error_timestamp.ToDatetime() >= current_dtime
+    assert error_timestamp.ToDatetime() >= current_timestamp.ToDatetime()
 
-    current_dtime = datetime.now()
+    current_timestamp = Timestamp()
+    current_timestamp.GetCurrentTime()
     error_timestamp = get_container_error_timestamp(FlyteException("foo"))
-    assert error_timestamp.ToDatetime() >= current_dtime
+    assert error_timestamp.ToDatetime() >= current_timestamp.ToDatetime()
 
-    current_dtime = datetime.now()
+    current_timestamp = Timestamp()
+    current_timestamp.GetCurrentTime()
     error_timestamp = get_container_error_timestamp(None)
-    assert error_timestamp.ToDatetime() >= current_dtime
+    assert error_timestamp.ToDatetime() >= current_timestamp.ToDatetime()
 
 
 def get_flyte_context(tmp_path_factory, outputs_path):