Fix security issue

Kaggle · Mar 30, 2024 · 1caa6b0 · 1caa6b0
1 parent 360163c
commit 1caa6b0
Showing 1 changed file with 6 additions and 9 deletions.
diff --git a/tests/test_kaggle_module_resolver.py b/tests/test_kaggle_module_resolver.py
@@ -64,9 +64,8 @@ def do_POST(self):
 
             # Load the files
             mount_slug = f"{model_ref['ModelSlug']}/{model_ref['Framework']}/{model_ref['InstanceSlug']}/{model_ref['VersionNumber']}"
-            model_path = os.path.join(MOUNT_PATH, mount_slug)
-            os.makedirs(os.path.dirname(model_path))
-            os.symlink('/input/tests/data/saved_model/', model_path, target_is_directory=True)
+            os.makedirs(os.path.dirname(os.path.join(MOUNT_PATH, mount_slug)))
+            os.symlink('/input/tests/data/saved_model/', os.path.join(MOUNT_PATH, mount_slug), target_is_directory=True)
 
             # Return the response
             self.wfile.write(bytes(json.dumps({
@@ -87,9 +86,8 @@ def test_kaggle_resolver_long_url_succeeds(self):
             layer = hub.KerasLayer(model_url)
             self.assertEqual([1, 1], layer(test_inputs).shape)
         # Delete the files that were created in KaggleJwtHandler's do_POST method
-        model_path = os.path.join(MOUNT_PATH, "foomodule/tensorflow2/barvar/2")
-        os.unlink(model_path)
-        os.rmdir(os.path.dirname(model_path))
+        os.unlink(os.path.join(MOUNT_PATH, "foomodule/tensorflow2/barvar/2"))
+        os.rmdir(os.path.dirname(os.path.join(MOUNT_PATH, "foomodule/tensorflow2/barvar/2")))
 
     def test_kaggle_resolver_short_url_succeeds(self):
         model_url = "https://kaggle.com/models/foo/foomodule/TensorFlow2/barvar/2"
@@ -98,9 +96,8 @@ def test_kaggle_resolver_short_url_succeeds(self):
             layer = hub.KerasLayer(model_url)
             self.assertEqual([1, 1], layer(test_inputs).shape)
         # Delete the files that were created in KaggleJwtHandler's do_POST method
-        model_path = os.path.join(MOUNT_PATH, "foomodule/tensorflow2/barvar/2")
-        os.unlink(model_path)
-        os.rmdir(os.path.dirname(model_path))
+        os.unlink(os.path.join(MOUNT_PATH, "foomodule/tensorflow2/barvar/2"))
+        os.rmdir(os.path.dirname(os.path.join(MOUNT_PATH, "foomodule/tensorflow2/barvar/2")))
 
     def test_kaggle_resolver_not_attached_throws(self):
         with create_test_server(KaggleJwtHandler) as addr: