1package dbtest
2
3import (
4	"bytes"
5	"fmt"
6	"net"
7	"os"
8	"os/exec"
9	"strconv"
10	"time"
11
12	"gopkg.in/mgo.v2"
13	"gopkg.in/tomb.v2"
14)
15
16// DBServer controls a MongoDB server process to be used within test suites.
17//
18// The test server is started when Session is called the first time and should
19// remain running for the duration of all tests, with the Wipe method being
20// called between tests (before each of them) to clear stored data. After all tests
21// are done, the Stop method should be called to stop the test server.
22//
23// Before the DBServer is used the SetPath method must be called to define
24// the location for the database files to be stored.
25type DBServer struct {
26	session *mgo.Session
27	output  bytes.Buffer
28	server  *exec.Cmd
29	dbpath  string
30	host    string
31	tomb    tomb.Tomb
32}
33
34// SetPath defines the path to the directory where the database files will be
35// stored if it is started. The directory path itself is not created or removed
36// by the test helper.
37func (dbs *DBServer) SetPath(dbpath string) {
38	dbs.dbpath = dbpath
39}
40
41func (dbs *DBServer) start() {
42	if dbs.server != nil {
43		panic("DBServer already started")
44	}
45	if dbs.dbpath == "" {
46		panic("DBServer.SetPath must be called before using the server")
47	}
48	mgo.SetStats(true)
49	l, err := net.Listen("tcp", "127.0.0.1:0")
50	if err != nil {
51		panic("unable to listen on a local address: " + err.Error())
52	}
53	addr := l.Addr().(*net.TCPAddr)
54	l.Close()
55	dbs.host = addr.String()
56
57	args := []string{
58		"--dbpath", dbs.dbpath,
59		"--bind_ip", "127.0.0.1",
60		"--port", strconv.Itoa(addr.Port),
61		"--nssize", "1",
62		"--noprealloc",
63		"--smallfiles",
64		"--nojournal",
65	}
66	dbs.tomb = tomb.Tomb{}
67	dbs.server = exec.Command("mongod", args...)
68	dbs.server.Stdout = &dbs.output
69	dbs.server.Stderr = &dbs.output
70	err = dbs.server.Start()
71	if err != nil {
72		panic(err)
73	}
74	dbs.tomb.Go(dbs.monitor)
75	dbs.Wipe()
76}
77
78func (dbs *DBServer) monitor() error {
79	dbs.server.Process.Wait()
80	if dbs.tomb.Alive() {
81		// Present some debugging information.
82		fmt.Fprintf(os.Stderr, "---- mongod process died unexpectedly:\n")
83		fmt.Fprintf(os.Stderr, "%s", dbs.output.Bytes())
84		fmt.Fprintf(os.Stderr, "---- mongod processes running right now:\n")
85		cmd := exec.Command("/bin/sh", "-c", "ps auxw | grep mongod")
86		cmd.Stdout = os.Stderr
87		cmd.Stderr = os.Stderr
88		cmd.Run()
89		fmt.Fprintf(os.Stderr, "----------------------------------------\n")
90
91		panic("mongod process died unexpectedly")
92	}
93	return nil
94}
95
96// Stop stops the test server process, if it is running.
97//
98// It's okay to call Stop multiple times. After the test server is
99// stopped it cannot be restarted.
100//
101// All database sessions must be closed before or while the Stop method
102// is running. Otherwise Stop will panic after a timeout informing that
103// there is a session leak.
104func (dbs *DBServer) Stop() {
105	if dbs.session != nil {
106		dbs.checkSessions()
107		if dbs.session != nil {
108			dbs.session.Close()
109			dbs.session = nil
110		}
111	}
112	if dbs.server != nil {
113		dbs.tomb.Kill(nil)
114		dbs.server.Process.Signal(os.Interrupt)
115		select {
116		case <-dbs.tomb.Dead():
117		case <-time.After(5 * time.Second):
118			panic("timeout waiting for mongod process to die")
119		}
120		dbs.server = nil
121	}
122}
123
124// Session returns a new session to the server. The returned session
125// must be closed after the test is done with it.
126//
127// The first Session obtained from a DBServer will start it.
128func (dbs *DBServer) Session() *mgo.Session {
129	if dbs.server == nil {
130		dbs.start()
131	}
132	if dbs.session == nil {
133		mgo.ResetStats()
134		var err error
135		dbs.session, err = mgo.Dial(dbs.host + "/test")
136		if err != nil {
137			panic(err)
138		}
139	}
140	return dbs.session.Copy()
141}
142
143// checkSessions ensures all mgo sessions opened were properly closed.
144// For slightly faster tests, it may be disabled setting the
145// environmnet variable CHECK_SESSIONS to 0.
146func (dbs *DBServer) checkSessions() {
147	if check := os.Getenv("CHECK_SESSIONS"); check == "0" || dbs.server == nil || dbs.session == nil {
148		return
149	}
150	dbs.session.Close()
151	dbs.session = nil
152	for i := 0; i < 100; i++ {
153		stats := mgo.GetStats()
154		if stats.SocketsInUse == 0 && stats.SocketsAlive == 0 {
155			return
156		}
157		time.Sleep(100 * time.Millisecond)
158	}
159	panic("There are mgo sessions still alive.")
160}
161
162// Wipe drops all created databases and their data.
163//
164// The MongoDB server remains running if it was prevoiusly running,
165// or stopped if it was previously stopped.
166//
167// All database sessions must be closed before or while the Wipe method
168// is running. Otherwise Wipe will panic after a timeout informing that
169// there is a session leak.
170func (dbs *DBServer) Wipe() {
171	if dbs.server == nil || dbs.session == nil {
172		return
173	}
174	dbs.checkSessions()
175	sessionUnset := dbs.session == nil
176	session := dbs.Session()
177	defer session.Close()
178	if sessionUnset {
179		dbs.session.Close()
180		dbs.session = nil
181	}
182	names, err := session.DatabaseNames()
183	if err != nil {
184		panic(err)
185	}
186	for _, name := range names {
187		switch name {
188		case "admin", "local", "config":
189		default:
190			err = session.DB(name).DropDatabase()
191			if err != nil {
192				panic(err)
193			}
194		}
195	}
196}
197