From 3b73e8c6435a8e42a011ba60cf746974236f0356 Mon Sep 17 00:00:00 2001
From: Raphael Vigee <raphvigee@gmail.com>
Date: Sat, 21 Nov 2020 03:54:48 +0100
Subject: [PATCH] Add tests & change count to int64

---
 README.md                     |  2 +-
 driver/driver.go              |  2 +-
 driver/gorm/page.go           |  4 ++--
 driver/gorm/paginator_test.go | 12 ++++++++++--
 4 files changed, 14 insertions(+), 6 deletions(-)

diff --git a/README.md b/README.md
index a9c7d5d..8400537 100644
--- a/README.md
+++ b/README.md
@@ -74,7 +74,7 @@ func main() {
         CreatedAt: base.Add(6 * time.Hour),
     })
 
-    // Define the pagination criterias
+    // Define the pagination criteria
     pg := paginator.New(paginator.Options{
         Driver: gorm.Driver{
             Columns: []*gorm.Column{
diff --git a/driver/driver.go b/driver/driver.go
index 896bc84..1c1a576 100644
--- a/driver/driver.go
+++ b/driver/driver.go
@@ -19,7 +19,7 @@ type Driver interface {
 
 type Executor interface {
 	Query(dst interface{}) error
-	Count() (uint64, error)
+	Count() (int64, error)
 }
 
 type PageInfo struct {
diff --git a/driver/gorm/page.go b/driver/gorm/page.go
index 81e7dc4..9c48c1e 100644
--- a/driver/gorm/page.go
+++ b/driver/gorm/page.go
@@ -25,7 +25,7 @@ func (g gormDriverPage) Query(dst interface{}) error {
 	return g.tx.Find(dst).Error
 }
 
-func (g gormDriverPage) Count() (uint64, error) {
+func (g gormDriverPage) Count() (int64, error) {
 	if g.tx == nil {
 		return 0, nil
 	}
@@ -33,7 +33,7 @@ func (g gormDriverPage) Count() (uint64, error) {
 	var c int64
 	err := g.tx.Count(&c).Error
 
-	return uint64(c), err
+	return c, err
 }
 
 func (g gormDriverPage) Info() driver.PageInfo {
diff --git a/driver/gorm/paginator_test.go b/driver/gorm/paginator_test.go
index 0d39cb2..5c981ef 100644
--- a/driver/gorm/paginator_test.go
+++ b/driver/gorm/paginator_test.go
@@ -120,6 +120,10 @@ func testPaginator(t *testing.T, columns []*Column, typ cursor.Type, limit int,
 		ec, _ := res.Cursor(int64(limit-1))
 		assert.Equal(t, ec, res.PageInfo.EndCursor)
 
+		c, err := res.Count()
+		assert.NoError(t, err)
+		assert.Equal(t, int64(len(s.names)), c)
+
 		var users []User
 		err = res.Query(&users)
 		assert.NoError(t, err)
@@ -153,10 +157,10 @@ func TestFactory_Empty(t *testing.T) {
 		Driver: Driver{Columns: simpleColumns},
 	})
 
-	c, err := pg.Cursor("", cursor.After, 2)
+	csr, err := pg.Cursor("", cursor.After, 2)
 	assert.NoError(t, err)
 
-	res, err := pg.Paginate(c, tx)
+	res, err := pg.Paginate(csr, tx)
 	assert.NoError(t, err)
 
 	assert.False(t, res.PageInfo.HasPreviousPage)
@@ -164,6 +168,10 @@ func TestFactory_Empty(t *testing.T) {
 	assert.Empty(t, res.PageInfo.StartCursor)
 	assert.Empty(t, res.PageInfo.EndCursor)
 
+	c, err := res.Count()
+	assert.NoError(t, err)
+	assert.Equal(t, int64(0), c)
+
 	var users []User
 	err = res.Query(&users)
 	assert.NoError(t, err)