clair/vendor/gopkg.in/mgo.v2/testserver/testserver.go

package testserver

import (
	"bytes"
	"fmt"
	"net"
	"os"
	"os/exec"
	"strconv"
	"time"

	"gopkg.in/mgo.v2"
	"gopkg.in/tomb.v2"
)

// TestServer controls a MongoDB server process to be used within test suites.
//
// The test server is started when Session is called the first time and should
// remain running for the duration of all tests, with the Wipe method being
// called between tests (before each of them) to clear stored data. After all tests
// are done, the Stop method should be called to stop the test server.
//
// Before the TestServer is used the SetPath method must be called to define
// the location for the database files to be stored.
type TestServer struct {
	session *mgo.Session
	output  bytes.Buffer
	server  *exec.Cmd
	dbpath  string
	host    string
	tomb    tomb.Tomb
}

// SetPath defines the path to the directory where the database files will be
// stored if it is started. The directory path itself is not created or removed
// by the test helper.
func (ts *TestServer) SetPath(dbpath string) {
	ts.dbpath = dbpath
}

func (ts *TestServer) start() {
	if ts.server != nil {
		panic("TestServer already started")
	}
	if ts.dbpath == "" {
		panic("TestServer.SetPath must be called before using the server")
	}
	mgo.SetStats(true)
	l, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		panic("unable to listen on a local address: " + err.Error())
	}
	addr := l.Addr().(*net.TCPAddr)
	l.Close()
	ts.host = addr.String()

	args := []string{
		"--dbpath", ts.dbpath,
		"--bind_ip", "127.0.0.1",
		"--port", strconv.Itoa(addr.Port),
		"--nssize", "1",
		"--noprealloc",
		"--smallfiles",
		"--nojournal",
	}
	ts.tomb = tomb.Tomb{}
	ts.server = exec.Command("mongod", args...)
	ts.server.Stdout = &ts.output
	ts.server.Stderr = &ts.output
	err = ts.server.Start()
	if err != nil {
		panic(err)
	}
	ts.tomb.Go(ts.monitor)
	ts.Wipe()
}

func (ts *TestServer) monitor() error {
	ts.server.Process.Wait()
	if ts.tomb.Alive() {
		// Present some debugging information.
		fmt.Fprintf(os.Stderr, "---- mongod process died unexpectedly:\n")
		fmt.Fprintf(os.Stderr, "%s", ts.output.Bytes())
		fmt.Fprintf(os.Stderr, "---- mongod processes running right now:\n")
		cmd := exec.Command("/bin/sh", "-c", "ps auxw | grep mongod")
		cmd.Stdout = os.Stderr
		cmd.Stderr = os.Stderr
		cmd.Run()
		fmt.Fprintf(os.Stderr, "----------------------------------------\n")

		panic("mongod process died unexpectedly")
	}
	return nil
}

// Stop stops the test server process, if it is running.
//
// It's okay to call Stop multiple times. After the test server is
// stopped it cannot be restarted.
//
// All database sessions must be closed before or while the Stop method
// is running. Otherwise Stop will panic after a timeout informing that
// there is a session leak.
func (ts *TestServer) Stop() {
	if ts.session != nil {
		ts.checkSessions()
		if ts.session != nil {
			ts.session.Close()
			ts.session = nil
		}
	}
	if ts.server != nil {
		ts.tomb.Kill(nil)
		ts.server.Process.Kill()
		select {
		case <-ts.tomb.Dead():
		case <-time.After(5 * time.Second):
			panic("timeout waiting for mongod process to die")
		}
		ts.server = nil
	}
}

// Session returns a new session to the server. The returned session
// must be closed after the test is done with it.
//
// The first Session obtained from a TestServer will start it.
func (ts *TestServer) Session() *mgo.Session {
	if ts.server == nil {
		ts.start()
	}
	if ts.session == nil {
		mgo.ResetStats()
		var err error
		ts.session, err = mgo.Dial(ts.host + "/test")
		if err != nil {
			panic(err)
		}
	}
	return ts.session.Copy()
}

// checkSessions ensures all mgo sessions opened were properly closed.
// For slightly faster tests, it may be disabled setting the
// environmnet variable CHECK_SESSIONS to 0.
func (ts *TestServer) checkSessions() {
	if check := os.Getenv("CHECK_SESSIONS"); check == "0" || ts.server == nil || ts.session == nil {
		return
	}
	ts.session.Close()
	ts.session = nil
	for i := 0; i < 100; i++ {
		stats := mgo.GetStats()
		if stats.SocketsInUse == 0 && stats.SocketsAlive == 0 {
			return
		}
		time.Sleep(100 * time.Millisecond)
	}
	panic("There are mgo sessions still alive.")
}

// Wipe drops all created databases and their data.
//
// The MongoDB server remains running if it was prevoiusly running,
// or stopped if it was previously stopped.
//
// All database sessions must be closed before or while the Wipe method
// is running. Otherwise Wipe will panic after a timeout informing that
// there is a session leak.
func (ts *TestServer) Wipe() {
	if ts.server == nil || ts.session == nil {
		return
	}
	ts.checkSessions()
	sessionUnset := ts.session == nil
	session := ts.Session()
	defer session.Close()
	if sessionUnset {
		ts.session.Close()
		ts.session = nil
	}
	names, err := session.DatabaseNames()
	if err != nil {
		panic(err)
	}
	for _, name := range names {
		switch name {
		case "admin", "local", "config":
		default:
			err = session.DB(name).DropDatabase()
			if err != nil {
				panic(err)
			}
		}
	}
}