From cb9c405f46118aef4cee1fa352bcc8551bc4cf1a Mon Sep 17 00:00:00 2001
From: Jordan <jmwright798@gmail.com>
Date: Wed, 26 Mar 2014 21:42:07 -0500
Subject: [PATCH] Added better testing with gocheck Fixed some typos leading to
 syntax errors TODO: Finish up gorm integration into templates

---
 config/config.go      |  2 --
 gophish.go            |  2 +-
 gophish_test.go       | 14 --------------
 models/campaign.go    |  6 +++---
 models/models.go      |  3 +--
 models/models_test.go | 38 ++++++++++++++++++++++++++++++++++++++
 6 files changed, 43 insertions(+), 22 deletions(-)
 delete mode 100644 gophish_test.go
 create mode 100644 models/models_test.go

diff --git a/config/config.go b/config/config.go
index 3884c9f7..a92e3a41 100644
--- a/config/config.go
+++ b/config/config.go
@@ -4,7 +4,6 @@ import (
 	"encoding/json"
 	"fmt"
 	"io/ioutil"
-	"os"
 )
 
 type SMTPServer struct {
@@ -27,7 +26,6 @@ func init() {
 	config_file, err := ioutil.ReadFile("./config.json")
 	if err != nil {
 		fmt.Printf("File error: %v\n", err)
-		os.Exit(1)
 	}
 	json.Unmarshal(config_file, &Conf)
 }
diff --git a/gophish.go b/gophish.go
index 9433e5ef..acc45e62 100644
--- a/gophish.go
+++ b/gophish.go
@@ -38,7 +38,7 @@ import (
 func main() {
 	//Setup the global variables and settings
 	err := models.Setup()
-	defer models.DB.Close()
+	//defer models.db.DB().Close()
 	if err != nil {
 		fmt.Println(err)
 	}
diff --git a/gophish_test.go b/gophish_test.go
deleted file mode 100644
index 97510b9a..00000000
--- a/gophish_test.go
+++ /dev/null
@@ -1,14 +0,0 @@
-package main
-
-import (
-	"testing"
-
-	"github.com/jordan-wright/gophish/models"
-)
-
-func TestDBSetup(t *testing.T) {
-	err := models.Setup()
-	if err != nil {
-		t.Fatalf("Failed creating database: %v", err)
-	}
-}
diff --git a/models/campaign.go b/models/campaign.go
index e3846b90..46d91887 100644
--- a/models/campaign.go
+++ b/models/campaign.go
@@ -33,7 +33,7 @@ func GetCampaigns(uid int64) ([]Campaign, error) {
 	if err != nil {
 		fmt.Println(err)
 	}
-	 */for i, _ := range cs {
+	for i, _ := range cs {
 		err := db.Model(&cs[i]).Related(&cs[i].Results).Error
 		if err != nil {
 			fmt.Println(err)
@@ -47,11 +47,11 @@ func GetCampaigns(uid int64) ([]Campaign, error) {
 func GetCampaign(id int64, uid int64) (Campaign, error) {
 	c := Campaign{}
 	err := db.Where("id = ?", id).Where("user_id = ?", uid).Find(&c).Error
-	 */if err != nil {
+	if err != nil {
 		return c, err
 	}
 	err = db.Model(&c).Related(&c.Results).Error
-	 */return c, err
+	return c, err
 }
 
 // PostCampaign inserts a campaign and all associated records into the database.
diff --git a/models/models.go b/models/models.go
index 24f98668..ecf572a3 100644
--- a/models/models.go
+++ b/models/models.go
@@ -1,7 +1,6 @@
 package models
 
 import (
-	"database/sql"
 	"errors"
 	"log"
 	"os"
@@ -14,7 +13,6 @@ import (
 
 var Conn *gorp.DbMap
 var db gorm.DB
-var DB *sql.DB
 var err error
 var ErrUsernameTaken = errors.New("username already taken")
 var Logger = log.New(os.Stdout, " ", log.Ldate|log.Ltime|log.Lshortfile)
@@ -29,6 +27,7 @@ type Flash struct {
 // It also populates the Gophish Config object
 func Setup() error {
 	db, err = gorm.Open("sqlite3", config.Conf.DBPath)
+	db.LogMode(true)
 	db.SetLogger(Logger)
 	if err != nil {
 		Logger.Println(err)
diff --git a/models/models_test.go b/models/models_test.go
new file mode 100644
index 00000000..6c2d6733
--- /dev/null
+++ b/models/models_test.go
@@ -0,0 +1,38 @@
+package models
+
+import (
+	"os"
+	"testing"
+
+	"github.com/jordan-wright/gophish/config"
+	"launchpad.net/gocheck"
+)
+
+// Hook up gocheck into the "go test" runner.
+func Test(t *testing.T) { gocheck.TestingT(t) }
+
+type ModelsSuite struct{}
+
+var _ = gocheck.Suite(&ModelsSuite{})
+
+func (s *ModelsSuite) SetUpSuite(c *gocheck.C) {
+	config.Conf.DBPath = "../gophish_test.db"
+	err := Setup()
+	if err != nil {
+		c.Fatalf("Failed creating database: %v", err)
+	}
+}
+
+func (s *ModelsSuite) TestGetUser(c *gocheck.C) {
+	u, err := GetUser(1)
+	c.Assert(err, gocheck.Equals, nil)
+	c.Assert(u.Username, gocheck.Equals, "admin")
+}
+
+func (s *ModelsSuite) TearDownSuite(c *gocheck.C) {
+	db.DB().Close()
+	err := os.Remove(config.Conf.DBPath)
+	if err != nil {
+		c.Fatalf("Failed deleting test database: %v", err)
+	}
+}