From f18ee0dedb9c77e799976f0cf01b477c1032099e Mon Sep 17 00:00:00 2001
From: gitlawr <lawrleegle@gmail.com>
Date: Sun, 7 Apr 2019 08:41:18 +0800
Subject: [PATCH] Do not wrap a wrapped schema store

Problem:
Pagination is broken when the schema store is wrapped twice

Solution:
Do not wrap a wrapped schema store
---
 store/wrapper/wrapper.go      |  4 +++
 store/wrapper/wrapper_test.go | 52 +++++++++++++++++++++++++++++++++++
 2 files changed, 56 insertions(+)
 create mode 100644 store/wrapper/wrapper_test.go

diff --git a/store/wrapper/wrapper.go b/store/wrapper/wrapper.go
index 24a9f7b2a..54d1783cf 100644
--- a/store/wrapper/wrapper.go
+++ b/store/wrapper/wrapper.go
@@ -7,6 +7,10 @@ import (
 )
 
 func Wrap(store types.Store) types.Store {
+	if _, ok := store.(*StoreWrapper); ok {
+		return store
+	}
+
 	return &StoreWrapper{
 		store: store,
 	}
diff --git a/store/wrapper/wrapper_test.go b/store/wrapper/wrapper_test.go
new file mode 100644
index 000000000..bacba5738
--- /dev/null
+++ b/store/wrapper/wrapper_test.go
@@ -0,0 +1,52 @@
+package wrapper
+
+import (
+	"testing"
+
+	"github.com/rancher/norman/api/handler"
+	"github.com/rancher/norman/parse"
+	"github.com/rancher/norman/store/empty"
+	"github.com/rancher/norman/types"
+	"github.com/stretchr/testify/assert"
+)
+
+type testStore struct {
+	empty.Store
+}
+
+func (t *testStore) List(apiContext *types.APIContext, schema *types.Schema, opt *types.QueryOptions) ([]map[string]interface{}, error) {
+	return []map[string]interface{}{{"1": "1"}, {"2": "2"}, {"3": "3"}}, nil
+}
+
+func TestWrap(t *testing.T) {
+	store := &testStore{}
+	limit := int64(1)
+	opt := &types.QueryOptions{
+		Pagination: &types.Pagination{
+			Limit: &limit,
+		},
+	}
+	apiContext := &types.APIContext{
+		SubContextAttributeProvider: &parse.DefaultSubContextAttributeProvider{},
+		QueryFilter:                 handler.QueryFilter,
+		Pagination:                  opt.Pagination,
+	}
+
+	wrapped := Wrap(store)
+	if _, err := wrapped.List(apiContext, &types.Schema{}, opt); err != nil {
+		t.Fatal(err)
+	}
+	assert.True(t, apiContext.Pagination.Partial)
+	assert.Equal(t, int64(3), *apiContext.Pagination.Total)
+	assert.Equal(t, int64(1), *apiContext.Pagination.Limit)
+
+	wrappedTwice := Wrap(wrapped)
+	apiContext.Pagination = opt.Pagination
+	if _, err := wrappedTwice.List(apiContext, &types.Schema{}, opt); err != nil {
+		t.Fatal(err)
+	}
+	assert.True(t, apiContext.Pagination.Partial)
+	assert.Equal(t, int64(3), *apiContext.Pagination.Total)
+	assert.Equal(t, int64(1), *apiContext.Pagination.Limit)
+
+}