From f62c7b6dc7df2f260b8ac1cc34d407f11ebc721a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Bu=C4=9Fra=20Gedik?= <bgedik@gmail.com>
Date: Thu, 14 Nov 2024 13:46:44 -0800
Subject: [PATCH] Correct stow listing function (#6013)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: Buğra Gedik < bgedik@gmail.com>
---
 .../ioutils/remote_file_output_reader_test.go        | 12 ++++++------
 flytestdlib/storage/stow_store.go                    |  4 ++--
 flytestdlib/storage/stow_store_test.go               |  3 ++-
 3 files changed, 10 insertions(+), 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
index 1cd7099f78..59c7f48307 100644
--- a/flyteplugins/go/tasks/pluginmachinery/ioutils/remote_file_output_reader_test.go
+++ b/flyteplugins/go/tasks/pluginmachinery/ioutils/remote_file_output_reader_test.go
@@ -172,17 +172,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)
 
@@ -231,9 +231,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
index e86a199bda..d7488e0695 100644
--- 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
index aec59051f3..719aa43130 100644
--- 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) {