From e145effd1795b59bfdfdc3e845f98d18af3a62ab Mon Sep 17 00:00:00 2001
From: Denis Kirillov <denis@nspcc.ru>
Date: Sun, 4 Jul 2021 23:25:53 +0300
Subject: [PATCH] [#129] Fix root objects listing

Signed-off-by: Denis Kirillov <denis@nspcc.ru>
---
 api/layer/layer.go      | 23 ++++++++++++---------
 api/layer/layer_test.go | 44 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 58 insertions(+), 9 deletions(-)
 create mode 100644 api/layer/layer_test.go

diff --git a/api/layer/layer.go b/api/layer/layer.go
index d245182c7..3b9f4f440 100644
--- a/api/layer/layer.go
+++ b/api/layer/layer.go
@@ -285,15 +285,7 @@ func (n *layer) ListObjects(ctx context.Context, p *ListObjectsParams) (*ListObj
 		result.NextMarker = result.Objects[len(result.Objects)-1].Name
 	}
 
-	index := 0
-	for _, oi := range result.Objects {
-		if isDir := uniqNames[oi.Name]; isDir {
-			result.Objects = append(result.Objects[:index], result.Objects[index+1:]...)
-			result.Prefixes = append(result.Prefixes, oi.Name)
-		} else {
-			index++
-		}
-	}
+	fillPrefixes(&result, uniqNames)
 	if needDirectoryAsKey {
 		res := []*ObjectInfo{{
 			Name:    p.Prefix,
@@ -306,6 +298,19 @@ func (n *layer) ListObjects(ctx context.Context, p *ListObjectsParams) (*ListObj
 	return &result, nil
 }
 
+func fillPrefixes(result *ListObjectsInfo, directories map[string]bool) {
+	index := 0
+	for range result.Objects {
+		name := result.Objects[index].Name
+		if isDir := directories[name]; isDir {
+			result.Objects = append(result.Objects[:index], result.Objects[index+1:]...)
+			result.Prefixes = append(result.Prefixes, name)
+		} else {
+			index++
+		}
+	}
+}
+
 // GetObject from storage.
 func (n *layer) GetObject(ctx context.Context, p *GetObjectParams) error {
 	var (
diff --git a/api/layer/layer_test.go b/api/layer/layer_test.go
new file mode 100644
index 000000000..c079e6f1c
--- /dev/null
+++ b/api/layer/layer_test.go
@@ -0,0 +1,44 @@
+package layer
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestFillingPrefixes(t *testing.T) {
+	cases := []struct {
+		name             string
+		list             *ListObjectsInfo
+		directories      map[string]bool
+		expectedPrefixes []string
+		expectedObjects  []*ObjectInfo
+	}{
+		{
+			name: "3 dirs",
+			list: &ListObjectsInfo{
+				Objects: []*ObjectInfo{{Name: "dir/"}, {Name: "dir2/"}, {Name: "dir3/"}},
+			},
+			directories:      map[string]bool{"dir/": true, "dir2/": true, "dir3/": true},
+			expectedPrefixes: []string{"dir/", "dir2/", "dir3/"},
+			expectedObjects:  []*ObjectInfo{},
+		},
+		{
+			name: "1 obj, 3 dirs",
+			list: &ListObjectsInfo{
+				Objects: []*ObjectInfo{{Name: "dir/"}, {Name: "dir2/"}, {Name: "dir3/"}, {Name: "obj"}},
+			},
+			directories:      map[string]bool{"dir/": true, "dir2/": true, "dir3/": true},
+			expectedPrefixes: []string{"dir/", "dir2/", "dir3/"},
+			expectedObjects:  []*ObjectInfo{{Name: "obj"}},
+		},
+	}
+
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			fillPrefixes(tc.list, tc.directories)
+			require.Equal(t, tc.expectedPrefixes, tc.list.Prefixes)
+			require.Equal(t, tc.expectedObjects, tc.list.Objects)
+		})
+	}
+}