update

fugue-project · Aug 20, 2024 · bceef07 · bceef07
1 parent c6accb2
commit bceef07
Showing 1 changed file with 7 additions and 5 deletions.
diff --git a/tests/tune/concepts/test_logger.py b/tests/tune/concepts/test_logger.py
@@ -9,8 +9,10 @@ def test_logger_context():
     m1 = MetricLogger()
     m2 = MetricLogger()
 
-    with set_current_metric_logger(m1):
-        assert get_current_metric_logger() == m1
-        with set_current_metric_logger(m2):
-            assert get_current_metric_logger() == m2
-        assert get_current_metric_logger() == m1
+    with set_current_metric_logger(m1) as mm1:
+        assert get_current_metric_logger() is m1
+        assert mm1 is m1
+        with set_current_metric_logger(m2) as mm2:
+            assert get_current_metric_logger() is m2
+            assert mm2 is m2
+        assert get_current_metric_logger() is m1