BUG: fix custom path test error (#2200)

Co-authored-by: wuzhaoxin <[email protected]>
xorbitsai · Aug 30, 2024 · 2f8c3ce · 2f8c3ce
1 parent 10094f9
commit 2f8c3ce
Showing 1 changed file with 13 additions and 14 deletions.
diff --git a/xinference/model/image/tests/test_stable_diffusion.py b/xinference/model/image/tests/test_stable_diffusion.py
@@ -287,20 +287,19 @@ def test_register_custom_image():
         unregister_image,
     )
 
-    tmp_dir = tempfile.mktemp()
-
-    model_spec = CustomImageModelFamilyV1(
-        model_family="stable_diffusion",
-        model_name="my-custom-image",
-        model_id="my-custom-image",
-        model_uri=os.path.abspath(tmp_dir),
-    )
-
-    register_image(model_spec, persist=False)
-    assert model_spec in get_user_defined_images()
-
-    unregister_image(model_spec.model_name, raise_error=False)
-    assert model_spec not in get_user_defined_images()
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        model_spec = CustomImageModelFamilyV1(
+            model_family="stable_diffusion",
+            model_name="my-custom-image",
+            model_id="my-custom-image",
+            model_uri=os.path.abspath(tmp_dir),
+        )
+
+        register_image(model_spec, persist=False)
+        assert model_spec in get_user_defined_images()
+
+        unregister_image(model_spec.model_name, raise_error=False)
+        assert model_spec not in get_user_defined_images()
 
 
 def test_persist_custom_image():