Merge pull request #324 from rcutmore/models-test-cleanup

Clean up between model tests

models/models_test.go

@@ -6,6 +6,7 @@ import (
 
 	"github.com/PuerkitoBio/goquery"
 	"github.com/gophish/gophish/config"
+	"github.com/jinzhu/gorm"
 	"gopkg.in/check.v1"
 )
 
@@ -25,12 +26,35 @@ func (s *ModelsSuite) SetUpSuite(c *check.C) {
 	}
 }
 
+func (s *ModelsSuite) TearDownTest(c *check.C) {
+	// Clear database tables between each test. If new tables are
+	// used in this test suite they will need to be cleaned up here.
+	db.Delete(Group{})
+	db.Delete(Target{})
+	db.Delete(GroupTarget{})
+	db.Delete(SMTP{})
+	db.Delete(Page{})
+
+	// Reset users table to default state.
+	db.Not("id", 1).Delete(User{})
+	db.Model(User{}).Update("username", "admin")
+}
+
 func (s *ModelsSuite) TestGetUser(c *check.C) {
 	u, err := GetUser(1)
 	c.Assert(err, check.Equals, nil)
 	c.Assert(u.Username, check.Equals, "admin")
 }
 
+func (s *ModelsSuite) TestPutUser(c *check.C) {
+	u, err := GetUser(1)
+	u.Username = "admin_changed"
+	err = PutUser(&u)
+	c.Assert(err, check.Equals, nil)
+	u, err = GetUser(1)
+	c.Assert(u.Username, check.Equals, "admin_changed")
+}
+
 func (s *ModelsSuite) TestGeneratedAPIKey(c *check.C) {
 	u, err := GetUser(1)
 	c.Assert(err, check.Equals, nil)
@@ -63,9 +87,82 @@ func (s *ModelsSuite) TestPostGroupNoTargets(c *check.C) {
 	c.Assert(err, check.Equals, ErrNoTargetsSpecified)
 }
 
+func (s *ModelsSuite) TestGetGroups(c *check.C) {
+	// Add groups.
+	PostGroup(&Group{
+		Name:    "Test Group 1",
+		Targets: []Target{Target{Email: "test1@example.com"}},
+		UserId:  1,
+	})
+	PostGroup(&Group{
+		Name:    "Test Group 2",
+		Targets: []Target{Target{Email: "test2@example.com"}},
+		UserId:  1,
+	})
+
+	// Get groups and test result.
+	groups, err := GetGroups(1)
+	c.Assert(err, check.Equals, nil)
+	c.Assert(len(groups), check.Equals, 2)
+	c.Assert(len(groups[0].Targets), check.Equals, 1)
+	c.Assert(len(groups[1].Targets), check.Equals, 1)
+	c.Assert(groups[0].Name, check.Equals, "Test Group 1")
+	c.Assert(groups[1].Name, check.Equals, "Test Group 2")
+	c.Assert(groups[0].Targets[0].Email, check.Equals, "test1@example.com")
+	c.Assert(groups[1].Targets[0].Email, check.Equals, "test2@example.com")
+}
+
+func (s *ModelsSuite) TestGetGroupsNoGroups(c *check.C) {
+	groups, err := GetGroups(1)
+	c.Assert(err, check.Equals, nil)
+	c.Assert(len(groups), check.Equals, 0)
+}
+
+func (s *ModelsSuite) TestGetGroup(c *check.C) {
+	// Add group.
+	PostGroup(&Group{
+		Name:    "Test Group",
+		Targets: []Target{Target{Email: "test@example.com"}},
+		UserId:  1,
+	})
+
+	// Get group and test result.
+	group, err := GetGroup(1, 1)
+	c.Assert(err, check.Equals, nil)
+	c.Assert(len(group.Targets), check.Equals, 1)
+	c.Assert(group.Name, check.Equals, "Test Group")
+	c.Assert(group.Targets[0].Email, check.Equals, "test@example.com")
+}
+
+func (s *ModelsSuite) TestGetGroupNoGroups(c *check.C) {
+	_, err := GetGroup(1, 1)
+	c.Assert(err, check.Equals, gorm.ErrRecordNotFound)
+}
+
+func (s *ModelsSuite) TestGetGroupByName(c *check.C) {
+	// Add group.
+	PostGroup(&Group{
+		Name:    "Test Group",
+		Targets: []Target{Target{Email: "test@example.com"}},
+		UserId:  1,
+	})
+
+	// Get group and test result.
+	group, err := GetGroupByName("Test Group", 1)
+	c.Assert(err, check.Equals, nil)
+	c.Assert(len(group.Targets), check.Equals, 1)
+	c.Assert(group.Name, check.Equals, "Test Group")
+	c.Assert(group.Targets[0].Email, check.Equals, "test@example.com")
+}
+
+func (s *ModelsSuite) TestGetGroupByNameNoGroups(c *check.C) {
+	_, err := GetGroupByName("Test Group", 1)
+	c.Assert(err, check.Equals, gorm.ErrRecordNotFound)
+}
+
 func (s *ModelsSuite) TestPutGroup(c *check.C) {
 	// Add test group.
-	group := Group{Name: "Put Test Group"}
+	group := Group{Name: "Test Group"}
 	group.Targets = []Target{
 		Target{Email: "test1@example.com", FirstName: "First", LastName: "Example"},
 		Target{Email: "test2@example.com", FirstName: "Second", LastName: "Example"},
@@ -74,14 +171,14 @@ func (s *ModelsSuite) TestPutGroup(c *check.C) {
 	PostGroup(&group)
 
 	// Update one of group's targets.
-	group.Targets[0].FirstName = "New"
+	group.Targets[0].FirstName = "Updated"
 	err := PutGroup(&group)
 	c.Assert(err, check.Equals, nil)
 
 	// Verify updated target information.
 	targets, _ := GetTargets(group.Id)
 	c.Assert(targets[0].Email, check.Equals, "test1@example.com")
-	c.Assert(targets[0].FirstName, check.Equals, "New")
+	c.Assert(targets[0].FirstName, check.Equals, "Updated")
 	c.Assert(targets[0].LastName, check.Equals, "Example")
 	c.Assert(targets[1].Email, check.Equals, "test2@example.com")
 	c.Assert(targets[1].FirstName, check.Equals, "Second")
@@ -90,10 +187,10 @@ func (s *ModelsSuite) TestPutGroup(c *check.C) {
 
 func (s *ModelsSuite) TestPutGroupEmptyAttribute(c *check.C) {
 	// Add test group.
-	group := Group{Name: "Put Empty Attribute Test Group"}
+	group := Group{Name: "Test Group"}
 	group.Targets = []Target{
-		Target{Email: "test3@example.com", FirstName: "Third", LastName: "Example"},
+		Target{Email: "test1@example.com", FirstName: "First", LastName: "Example"},
-		Target{Email: "test4@example.com", FirstName: "Fourth", LastName: "Example"},
+		Target{Email: "test2@example.com", FirstName: "Second", LastName: "Example"},
 	}
 	group.UserId = 1
 	PostGroup(&group)
@@ -105,11 +202,11 @@ func (s *ModelsSuite) TestPutGroupEmptyAttribute(c *check.C) {
 
 	// Verify updated empty attribute was saved.
 	targets, _ := GetTargets(group.Id)
-	c.Assert(targets[0].Email, check.Equals, "test3@example.com")
+	c.Assert(targets[0].Email, check.Equals, "test1@example.com")
 	c.Assert(targets[0].FirstName, check.Equals, "")
 	c.Assert(targets[0].LastName, check.Equals, "Example")
-	c.Assert(targets[1].Email, check.Equals, "test4@example.com")
+	c.Assert(targets[1].Email, check.Equals, "test2@example.com")
-	c.Assert(targets[1].FirstName, check.Equals, "Fourth")
+	c.Assert(targets[1].FirstName, check.Equals, "Second")
 	c.Assert(targets[1].LastName, check.Equals, "Example")
 }
 
@@ -223,12 +320,3 @@ func (s *ModelsSuite) TestPostPage(c *check.C) {
 		c.Assert(ok, check.Equals, false)
 	})
 }
-
-func (s *ModelsSuite) TestPutUser(c *check.C) {
-	u, err := GetUser(1)
-	u.Username = "admin_changed"
-	err = PutUser(&u)
-	c.Assert(err, check.Equals, nil)
-	u, err = GetUser(1)
-	c.Assert(u.Username, check.Equals, "admin_changed")
-}