Skip to content

Commit 1caa6b0

Browse files
committed
Fix security issue
1 parent 360163c commit 1caa6b0

File tree

1 file changed

+6
-9
lines changed

1 file changed

+6
-9
lines changed

tests/test_kaggle_module_resolver.py

Lines changed: 6 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -64,9 +64,8 @@ def do_POST(self):
6464

6565
# Load the files
6666
mount_slug = f"{model_ref['ModelSlug']}/{model_ref['Framework']}/{model_ref['InstanceSlug']}/{model_ref['VersionNumber']}"
67-
model_path = os.path.join(MOUNT_PATH, mount_slug)
68-
os.makedirs(os.path.dirname(model_path))
69-
os.symlink('/input/tests/data/saved_model/', model_path, target_is_directory=True)
67+
os.makedirs(os.path.dirname(os.path.join(MOUNT_PATH, mount_slug)))
68+
os.symlink('/input/tests/data/saved_model/', os.path.join(MOUNT_PATH, mount_slug), target_is_directory=True)
7069

7170
# Return the response
7271
self.wfile.write(bytes(json.dumps({
@@ -87,9 +86,8 @@ def test_kaggle_resolver_long_url_succeeds(self):
8786
layer = hub.KerasLayer(model_url)
8887
self.assertEqual([1, 1], layer(test_inputs).shape)
8988
# Delete the files that were created in KaggleJwtHandler's do_POST method
90-
model_path = os.path.join(MOUNT_PATH, "foomodule/tensorflow2/barvar/2")
91-
os.unlink(model_path)
92-
os.rmdir(os.path.dirname(model_path))
89+
os.unlink(os.path.join(MOUNT_PATH, "foomodule/tensorflow2/barvar/2"))
90+
os.rmdir(os.path.dirname(os.path.join(MOUNT_PATH, "foomodule/tensorflow2/barvar/2")))
9391

9492
def test_kaggle_resolver_short_url_succeeds(self):
9593
model_url = "https://kaggle.com/models/foo/foomodule/TensorFlow2/barvar/2"
@@ -98,9 +96,8 @@ def test_kaggle_resolver_short_url_succeeds(self):
9896
layer = hub.KerasLayer(model_url)
9997
self.assertEqual([1, 1], layer(test_inputs).shape)
10098
# Delete the files that were created in KaggleJwtHandler's do_POST method
101-
model_path = os.path.join(MOUNT_PATH, "foomodule/tensorflow2/barvar/2")
102-
os.unlink(model_path)
103-
os.rmdir(os.path.dirname(model_path))
99+
os.unlink(os.path.join(MOUNT_PATH, "foomodule/tensorflow2/barvar/2"))
100+
os.rmdir(os.path.dirname(os.path.join(MOUNT_PATH, "foomodule/tensorflow2/barvar/2")))
104101

105102
def test_kaggle_resolver_not_attached_throws(self):
106103
with create_test_server(KaggleJwtHandler) as addr:

0 commit comments

Comments
 (0)