diff --git a/tests/tune/concepts/test_logger.py b/tests/tune/concepts/test_logger.py
index 7187acf..3dd4d2b 100644
--- 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