1 files changed, 15 insertions, 7 deletions
diff --git a/workhorse/internal/upload/uploads_test.go b/workhorse/internal/upload/uploads_test.go
index 0885f31d5a4..d77d73b5f48 100644
--- a/workhorse/internal/upload/uploads_test.go
+++ b/workhorse/internal/upload/uploads_test.go
@@ -325,14 +325,20 @@ func TestInvalidFileNames(t *testing.T) {
 	defer os.RemoveAll(tempPath)
 
 	for _, testCase := range []struct {
-		filename string
-		code     int
+		filename               string
+		code                   int
+		FeatureFlagExtractBase bool
+		expectedPrefix         string
 	}{
-		{"foobar", 200}, // sanity check for test setup below
-		{"foo/bar", 500},
-		{"/../../foobar", 500},
-		{".", 500},
-		{"..", 500},
+		{"foobar", 200, false, "foobar"}, // sanity check for test setup below
+		{"foo/bar", 500, false, ""},
+		{"foo/bar", 200, true, "bar"},
+		{"foo/bar/baz", 200, true, "baz"},
+		{"/../../foobar", 500, false, ""},
+		{"/../../foobar", 200, true, "foobar"},
+		{".", 500, false, ""},
+		{"..", 500, false, ""},
+		{"./", 500, false, ""},
 	} {
 		buffer := &bytes.Buffer{}
 
@@ -350,10 +356,12 @@ func TestInvalidFileNames(t *testing.T) {
 		apiResponse := &api.Response{TempPath: tempPath}
 		preparer := &DefaultPreparer{}
 		opts, _, err := preparer.Prepare(apiResponse)
+		opts.FeatureFlagExtractBase = testCase.FeatureFlagExtractBase
 		require.NoError(t, err)
 
 		HandleFileUploads(response, httpRequest, nilHandler, apiResponse, &SavedFileTracker{Request: httpRequest}, opts)
 		require.Equal(t, testCase.code, response.Code)
+		require.Equal(t, testCase.expectedPrefix, opts.TempFilePrefix)
 	}
 }