gophish/models/models_test.go

140 lines
3.6 KiB
Go
Raw Normal View History

package models
import (
"testing"
"github.com/gophish/gophish/config"
"gopkg.in/check.v1"
)
// Hook up gocheck into the "go test" runner.
func Test(t *testing.T) { check.TestingT(t) }
type ModelsSuite struct {
config *config.Config
}
var _ = check.Suite(&ModelsSuite{})
func (s *ModelsSuite) SetUpSuite(c *check.C) {
conf := &config.Config{
DBName: "sqlite3",
DBPath: ":memory:",
MigrationsPath: "../db/db_sqlite3/migrations/",
}
s.config = conf
err := Setup(conf)
if err != nil {
c.Fatalf("Failed creating database: %v", err)
}
}
2016-07-15 11:28:11 +00:00
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{})
db.Delete(Result{})
db.Delete(MailLog{})
db.Delete(Campaign{})
2016-07-15 11:28:11 +00:00
// Reset users table to default state.
db.Not("id", 1).Delete(User{})
db.Model(User{}).Update("username", "admin")
}
func (s *ModelsSuite) createCampaignDependencies(ch *check.C, optional ...string) Campaign {
// we use the optional parameter to pass an alternative subject
group := Group{Name: "Test Group"}
group.Targets = []Target{
Target{BaseRecipient: BaseRecipient{Email: "test1@example.com", FirstName: "First", LastName: "Example"}},
Target{BaseRecipient: BaseRecipient{Email: "test2@example.com", FirstName: "Second", LastName: "Example"}},
Target{BaseRecipient: BaseRecipient{Email: "test3@example.com", FirstName: "Second", LastName: "Example"}},
Target{BaseRecipient: BaseRecipient{Email: "test4@example.com", FirstName: "Second", LastName: "Example"}},
}
group.UserId = 1
ch.Assert(PostGroup(&group), check.Equals, nil)
// Add a template
t := Template{Name: "Test Template"}
if len(optional) > 0 {
t.Subject = optional[0]
} else {
t.Subject = "{{.RId}} - Subject"
}
t.Text = "{{.RId}} - Text"
t.HTML = "{{.RId}} - HTML"
t.UserId = 1
ch.Assert(PostTemplate(&t), check.Equals, nil)
// Add a landing page
p := Page{Name: "Test Page"}
p.HTML = "<html>Test</html>"
p.UserId = 1
ch.Assert(PostPage(&p), check.Equals, nil)
// Add a sending profile
smtp := SMTP{Name: "Test Page"}
smtp.UserId = 1
smtp.Host = "example.com"
smtp.FromAddress = "test@test.com"
ch.Assert(PostSMTP(&smtp), check.Equals, nil)
c := Campaign{Name: "Test campaign"}
c.UserId = 1
c.Template = t
c.Page = p
c.SMTP = smtp
c.Groups = []Group{group}
return c
}
func (s *ModelsSuite) createCampaign(ch *check.C) Campaign {
c := s.createCampaignDependencies(ch)
// Setup and "launch" our campaign
ch.Assert(PostCampaign(&c, c.UserId), check.Equals, nil)
// For comparing the dates, we need to fetch the campaign again. This is
// to solve an issue where the campaign object right now has time down to
// the microsecond, while in MySQL it's rounded down to the second.
c, _ = GetCampaign(c.Id, c.UserId)
return c
}
func setupBenchmark(b *testing.B) {
conf := &config.Config{
DBName: "sqlite3",
DBPath: ":memory:",
MigrationsPath: "../db/db_sqlite3/migrations/",
}
err := Setup(conf)
if err != nil {
b.Fatalf("Failed creating database: %v", err)
}
}
func tearDownBenchmark(b *testing.B) {
err := db.Close()
if err != nil {
b.Fatalf("error closing database: %v", err)
}
}
func resetBenchmark(b *testing.B) {
db.Delete(Group{})
db.Delete(Target{})
db.Delete(GroupTarget{})
db.Delete(SMTP{})
db.Delete(Page{})
db.Delete(Result{})
db.Delete(MailLog{})
db.Delete(Campaign{})
// Reset users table to default state.
db.Not("id", 1).Delete(User{})
db.Model(User{}).Update("username", "admin")
}