From 36170724dd658b453739a225af8b01d437585bc4 Mon Sep 17 00:00:00 2001
From: Ian Tan <ian.tan@techinasia.com>
Date: Fri, 17 Nov 2017 12:29:06 +0800
Subject: [PATCH] Refactor mock connection open/close logic

---
 test_helper.go         | 22 +++++++++++++++-------
 test_helper_adapter.go | 18 ++++++++----------
 test_helper_test.go    | 29 +++++++++++++++++++----------
 3 files changed, 42 insertions(+), 27 deletions(-)

diff --git a/test_helper.go b/test_helper.go
index a6203554..d9b99118 100644
--- a/test_helper.go
+++ b/test_helper.go
@@ -1,14 +1,11 @@
 package gorm
 
 import (
-	"database/sql"
 	"fmt"
 )
 
 // TestHelper is the exported struct used for setting expectations
 type TestHelper struct {
-	gormDb   *DB
-	mockDb   *sql.DB
 	adapter  Adapter
 	asserter Asserter
 }
@@ -29,24 +26,35 @@ func (h *TestHelper) ExpectFind(model interface{}) {
 
 // NewTestHelper returns a fresh TestHelper with an arbitary Adapter
 func NewTestHelper(adapter Adapter) (error, *DB, *TestHelper) {
-	err, mockDb, gormDb, asserter := adapter.Open()
+	err, asserter := adapter.Open()
 
 	if err != nil {
 		return err, nil, nil
 	}
 
-	return nil, gormDb, &TestHelper{gormDb: gormDb, mockDb: mockDb, adapter: adapter, asserter: asserter}
+	gormDb, err := Open("sqlmock", "mock_gorm_dsn")
+
+	if err != nil {
+		return err, nil, nil
+	}
+
+	return nil, gormDb, &TestHelper{adapter: adapter, asserter: asserter}
 }
 
 // NewDefaultTestHelper returns a TestHelper powered by go-sqlmock
 func NewDefaultTestHelper() (error, *DB, *TestHelper) {
 	adapter := &SqlmockAdapter{}
-	err, mockDb, gormDb, asserter := adapter.Open()
+	err, asserter := adapter.Open()
 
 	if err != nil {
 		return err, nil, nil
 	}
 
-	return nil, gormDb, &TestHelper{gormDb: gormDb, mockDb: mockDb, adapter: adapter, asserter: asserter}
+	gormDb, err := Open("sqlmock", "mock_gorm_dsn")
 
+	if err != nil {
+		return err, nil, nil
+	}
+
+	return nil, gormDb, &TestHelper{adapter: adapter, asserter: asserter}
 }
diff --git a/test_helper_adapter.go b/test_helper_adapter.go
index 595edf22..260cd6e3 100644
--- a/test_helper_adapter.go
+++ b/test_helper_adapter.go
@@ -7,30 +7,28 @@ import (
 )
 
 type Adapter interface {
-	Open() (error, *sql.DB, *DB, Asserter)
+	Open() (error, Asserter)
 	Close() error
 }
 
+// SqlmockAdapter implemenets the Adapter interface using go-sqlmock
+// it is the default Adapter
 type SqlmockAdapter struct {
 	mockDb *sql.DB
 	mock   *sqlmock.Sqlmock
 }
 
-// Open returns the raw sql.DB and a gorm DB instance
-func (adapter *SqlmockAdapter) Open() (error, *sql.DB, *DB, Asserter) {
+// Open returns the raw sql.DB instance and an Asserter
+func (adapter *SqlmockAdapter) Open() (error, Asserter) {
 	mockDb, mock, err := sqlmock.NewWithDSN("mock_gorm_dsn")
 
-	if err != nil {
-		return err, nil, nil, nil
-	}
-
-	gormDb, err := Open("sqlmock", "mock_gorm_dsn")
+	adapter.mockDb = mockDb
 
 	if err != nil {
-		return err, nil, nil, nil
+		return err, nil
 	}
 
-	return nil, mockDb, gormDb, &SqlmockAsserter{mock: mock, sqlmockDB: mockDb}
+	return nil, &SqlmockAsserter{mock: mock, sqlmockDB: mockDb}
 }
 
 func (adapter *SqlmockAdapter) Close() error {
diff --git a/test_helper_test.go b/test_helper_test.go
index 39d3f0c1..f597933b 100644
--- a/test_helper_test.go
+++ b/test_helper_test.go
@@ -6,21 +6,30 @@ import (
 	"github.com/jinzhu/gorm"
 )
 
-var helper *gorm.TestHelper
-var mockDb *gorm.DB
+func TestHelperOpen(t *testing.T) {
+	err, helper, db := gorm.NewTestHelper(&gorm.SqlmockAdapter{})
 
-func TestOpenWithSqlmock(t *testing.T) {
-	err, db, h := gorm.NewTestHelper(&gorm.SqlmockAdapter{})
+	defer func() {
+		helper.Close()
+		db.Close()
+	}()
 
 	if err != nil {
 		t.Fatal(err)
 	}
-
-	helper = h
-	mockDb = db
 }
 
-func TestQuery(t *testing.T) {
-	helper.ExpectFirst(&User{}).Return(&User{})
-	mockDb.First(&User{})
+func TestHelperClose(t *testing.T) {
+	err, helper, _ := gorm.NewTestHelper(&gorm.SqlmockAdapter{})
+
+	closeErr := helper.Close()
+
+	if err != nil {
+		t.Fatal(closeErr)
+	}
 }
+
+// func TestQuery(t *testing.T) {
+// 	helper.ExpectFirst(&User{}).Return(&User{})
+// 	mockDb.First(&User{})
+// }