Correct stow listing function (#6013)

Signed-off-by: Buğra Gedik < [email protected]>
flyteorg · Nov 14, 2024 · f160c52 · f160c52
1 parent afad1b2
commit f160c52
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 9 deletions.
diff --git a/flyteplugins/go/tasks/pluginmachinery/ioutils/remote_file_output_reader_test.go b/flyteplugins/go/tasks/pluginmachinery/ioutils/remote_file_output_reader_test.go
@@ -145,17 +145,17 @@ func TestReadOrigin(t *testing.T) {
 		}).Return(nil)
 
 		store.OnList(ctx, storage.DataReference("s3://errors/error"), 1000, storage.NewCursorAtStart()).Return(
-			[]storage.DataReference{"error-0.pb", "error-1.pb", "error-2.pb"}, storage.NewCursorAtEnd(), nil)
+			[]storage.DataReference{"s3://errors/error-0.pb", "s3://errors/error-1.pb", "s3://errors/error-2.pb"}, storage.NewCursorAtEnd(), nil)
 
-		store.OnHead(ctx, storage.DataReference("error-0.pb")).Return(MemoryMetadata{
+		store.OnHead(ctx, storage.DataReference("s3://errors/error-0.pb")).Return(MemoryMetadata{
 			exists: true,
 		}, nil)
 
-		store.OnHead(ctx, storage.DataReference("error-1.pb")).Return(MemoryMetadata{
+		store.OnHead(ctx, storage.DataReference("s3://errors/error-1.pb")).Return(MemoryMetadata{
 			exists: true,
 		}, nil)
 
-		store.OnHead(ctx, storage.DataReference("error-2.pb")).Return(MemoryMetadata{
+		store.OnHead(ctx, storage.DataReference("s3://errors/error-2.pb")).Return(MemoryMetadata{
 			exists: true,
 		}, nil)
 
@@ -204,9 +204,9 @@ func TestReadOrigin(t *testing.T) {
 		}).Return(nil)
 
 		store.OnList(ctx, storage.DataReference("s3://errors/error"), 1000, storage.NewCursorAtStart()).Return(
-			[]storage.DataReference{"error.pb"}, storage.NewCursorAtEnd(), nil)
+			[]storage.DataReference{"s3://errors/error.pb"}, storage.NewCursorAtEnd(), nil)
 
-		store.OnHead(ctx, storage.DataReference("error.pb")).Return(MemoryMetadata{
+		store.OnHead(ctx, storage.DataReference("s3://errors/error.pb")).Return(MemoryMetadata{
 			exists: true,
 		}, nil)
 

diff --git a/flytestdlib/storage/stow_store.go b/flytestdlib/storage/stow_store.go
@@ -263,7 +263,7 @@ func (s *StowStore) Head(ctx context.Context, reference DataReference) (Metadata
 }
 
 func (s *StowStore) List(ctx context.Context, reference DataReference, maxItems int, cursor Cursor) ([]DataReference, Cursor, error) {
-	scheme, containerName, key, err := reference.Split()
+	_, containerName, key, err := reference.Split()
 	if err != nil {
 		s.metrics.BadReference.Inc(ctx)
 		return nil, NewCursorAtEnd(), err
@@ -291,7 +291,7 @@ func (s *StowStore) List(ctx context.Context, reference DataReference, maxItems
 	if err == nil {
 		results := make([]DataReference, len(items))
 		for index, item := range items {
-			results[index] = DataReference(fmt.Sprintf("%s://%s/%s", scheme, containerName, item.URL().String()))
+			results[index] = DataReference(item.URL().String())
 		}
 		if stow.IsCursorEnd(stowCursor) {
 			cursor = NewCursorAtEnd()

diff --git a/flytestdlib/storage/stow_store_test.go b/flytestdlib/storage/stow_store_test.go
@@ -97,7 +97,8 @@ func (m mockStowContainer) Items(prefix, cursor string, count int) ([]stow.Item,
 	numItems := endIndexExc - startIndex
 	results := make([]stow.Item, numItems)
 	for index, itemKey := range itemKeys[startIndex:endIndexExc] {
-		results[index] = m.items[itemKey]
+		url := fmt.Sprintf("s3://%s/%s", m.id, m.items[itemKey].url)
+		results[index] = mockStowItem{url: url, size: m.items[itemKey].size}
 	}
 
 	if endIndexExc == len(m.items) {