From 8bac8e9a36ebefc12b470542aad6503d7370e482 Mon Sep 17 00:00:00 2001
From: Delio Vicini <vicini@google.com>
Date: Mon, 12 Aug 2024 15:29:47 +0200
Subject: [PATCH] Fix Python ScopedThreadEnvironment

---
 src/core/python/thread.cpp    | 3 ++-
 src/core/tests/test_thread.py | 7 +++++++
 2 files changed, 9 insertions(+), 1 deletion(-)
 create mode 100644 src/core/tests/test_thread.py

diff --git a/src/core/python/thread.cpp b/src/core/python/thread.cpp
index 09dea61c8..e76d381b0 100644
--- a/src/core/python/thread.cpp
+++ b/src/core/python/thread.cpp
@@ -112,6 +112,7 @@ MI_PY_EXPORT(Thread) {
                                             D(ScopedSetThreadEnvironment))
         .def(nb::init<const ThreadEnvironment &>())
         .def("__enter__", &PyScopedSetThreadEnvironment::enter)
-        .def("__exit__", &PyScopedSetThreadEnvironment::exit);
+        .def("__exit__", &PyScopedSetThreadEnvironment::exit,
+             "exc_type"_a.none(), "exc_val"_a.none(), "exc_tb"_a.none());
 }
 
diff --git a/src/core/tests/test_thread.py b/src/core/tests/test_thread.py
new file mode 100644
index 000000000..96acef1b1
--- /dev/null
+++ b/src/core/tests/test_thread.py
@@ -0,0 +1,7 @@
+import pytest
+import mitsuba as mi
+
+def test01_use_scoped_thread_environment(variant_scalar_rgb):
+    environment = mi.ThreadEnvironment()
+    with mi.ScopedSetThreadEnvironment(environment):
+        mi.Log(mi.LogLevel.Info, "Log from a thread environment.")