full testing support for pilosa server

1. fixed a couple bugs with config file reading.
2. made cmd.Serve a global variable so tests could inspect it.
3. added table style testing for pilosa server
4. added the ability to stop server programmatically which also causes the cobra
command running it to return.

cmd/root.go

@@ -3,7 +3,6 @@ package cmd
 import (
 	"fmt"
 	"io"
-	"log"
 	"strings"
 
 	"github.com/spf13/cobra"
@@ -86,7 +85,8 @@ func setAllConfig(v *viper.Viper, flags *flag.FlagSet, envPrefix string) error {
 
 	// add config file to viper
 	if c != "" {
-		v.AddConfigPath(c)
+		v.SetConfigFile(c)
+		v.SetConfigType("toml")
 		err := v.ReadInConfig()
 		if err != nil {
 			return fmt.Errorf("error reading configuration file '%s': %v", c, err)
@@ -99,13 +99,8 @@ func setAllConfig(v *viper.Viper, flags *flag.FlagSet, envPrefix string) error {
 		if flagErr != nil {
 			return
 		}
-		log.Printf("Now visiting: %v with value '%s'", f.Name, f.Value)
 		value := v.GetString(f.Name)
-		log.Printf("Setting to value: '%v'", value)
 		flagErr = f.Value.Set(value)
 	})
-	if flagErr == nil {
-		fmt.Println(v.AllSettings())
-	}
 	return flagErr
 }

cmd/root_test.go

@@ -14,6 +14,16 @@ import (
 	"github.com/spf13/cobra"
 )
 
+func failErr(t *testing.T, err error, context ...string) {
+	ctx := strings.Join(context, "; ")
+	if err != nil {
+		t.Fatal(ctx, ": ", err)
+	}
+}
+
+// tExec executes the given `cmd`, which will be writing its output to `w`, and
+// can be read from `out`. It will fail the test if the command does not return
+// within 1 second. Useful for testing help messages and such.
 func tExec(t *testing.T, cmd *cobra.Command, out io.Reader, w io.WriteCloser) (output []byte) {
 	done := make(chan struct{})
 	go func() {
@@ -40,6 +50,9 @@ func tExec(t *testing.T, cmd *cobra.Command, out io.Reader, w io.WriteCloser) (o
 	return output
 }
 
+// ExecNewRootCommand executes the pilosa root command with the given arguments
+// and returns it's output. It will fail if the command does not complete within
+// 1 second.
 func ExecNewRootCommand(t *testing.T, args ...string) string {
 	out, w := io.Pipe()
 	rc := cmd.NewRootCommand(os.Stdin, w, w)

cmd/server.go

@@ -13,9 +13,12 @@ import (
 	"github.com/pilosa/pilosa/server"
 )
 
+// Serve is global so that tests can control and verify it.
+var Serve *server.Command
+
 func NewServeCmd(stdin io.Reader, stdout, stderr io.Writer) *cobra.Command {
-	serve := server.NewCommand()
-	serve.Stdin, serve.Stdout, serve.Stderr = stdin, stdout, stderr
+	Serve = server.NewCommand()
+	Serve.Stdin, Serve.Stdout, Serve.Stderr = stdin, stdout, stderr
 	serveCmd := &cobra.Command{
 		Use:   "server",
 		Short: "Run Pilosa.",
@@ -25,52 +28,57 @@ It will load existing data from the configured
 directory, and start listening client connections
 on the configured port.`,
 		RunE: func(cmd *cobra.Command, args []string) error {
-			serve.Server.Handler.Version = Version
-			fmt.Fprintf(serve.Stderr, "Pilosa %s, build time %s\n", Version, BuildTime)
+			Serve.Server.Handler.Version = Version
+			fmt.Fprintf(Serve.Stderr, "Pilosa %s, build time %s\n", Version, BuildTime)
 
 			// Start CPU profiling.
-			if serve.CPUProfile != "" {
-				f, err := os.Create(serve.CPUProfile)
+			if Serve.CPUProfile != "" {
+				f, err := os.Create(Serve.CPUProfile)
 				if err != nil {
 					return fmt.Errorf("create cpu profile: %v", err)
 				}
 				defer f.Close()
 
-				fmt.Fprintln(serve.Stderr, "Starting cpu profile")
+				fmt.Fprintln(Serve.Stderr, "Starting cpu profile")
 				pprof.StartCPUProfile(f)
-				time.AfterFunc(serve.CPUTime, func() {
-					fmt.Fprintln(serve.Stderr, "Stopping cpu profile")
+				time.AfterFunc(Serve.CPUTime, func() {
+					fmt.Fprintln(Serve.Stderr, "Stopping cpu profile")
 					pprof.StopCPUProfile()
 					f.Close()
 				})
 			}
 
 			// Execute the program.
-			if err := serve.Run(); err != nil {
+			if err := Serve.Run(); err != nil {
 				return err
 			}
 
 			// First SIGKILL causes server to shut down gracefully.
 			c := make(chan os.Signal, 2)
 			signal.Notify(c, os.Interrupt)
-			sig := <-c
-			fmt.Fprintf(serve.Stderr, "Received %s; gracefully shutting down...\n", sig.String())
+			select {
+			case sig := <-c:
+				fmt.Fprintf(Serve.Stderr, "Received %s; gracefully shutting down...\n", sig.String())
 
-			// Second signal causes a hard shutdown.
-			go func() { <-c; os.Exit(1) }()
+				// Second signal causes a hard shutdown.
+				go func() { <-c; os.Exit(1) }()
 
-			if err := serve.Close(); err != nil {
-				return err
+				if err := Serve.Close(); err != nil {
+					return err
+				}
+			case <-Serve.Done:
+				fmt.Fprintf(Serve.Stderr, "Server closed externally")
 			}
 			return nil
 		},
 	}
 	flags := serveCmd.Flags()
 
-	flags.StringVarP(&serve.ConfigPath, "config", "c", "", "Configuration file to read from.")
-	flags.StringVarP(&serve.Config.DataDir, "data-dir", "d", "~/.pilosa", "Directory to store pilosa data files.")
-	flags.StringVarP(&serve.CPUProfile, "cpu-profile", "", "", "Where to store CPU profile.")
-	flags.DurationVarP(&serve.CPUTime, "cpu-time", "", 30*time.Second, "CPU profile duration.")
+	flags.StringVarP(&Serve.ConfigPath, "config", "c", "", "Configuration file to read from.")
+	flags.StringVarP(&Serve.Config.Host, "host", "", ":10101", "Default URI on which pilosa should listen.")
+	flags.StringVarP(&Serve.Config.DataDir, "data-dir", "d", "~/.pilosa", "Directory to store pilosa data files.")
+	flags.StringVarP(&Serve.CPUProfile, "cpu-profile", "", "", "Where to store CPU profile.")
+	flags.DurationVarP(&Serve.CPUTime, "cpu-time", "", 30*time.Second, "CPU profile duration.")
 
 	return serveCmd
 }

cmd/server_test.go

@@ -1,8 +1,16 @@
 package cmd_test
 
 import (
+	"fmt"
+	"io/ioutil"
 	"strings"
+	"sync"
 	"testing"
+
+	"os"
+
+	"github.com/pilosa/pilosa/cmd"
+	"github.com/spf13/cobra"
 )
 
 func TestServerHelp(t *testing.T) {
@@ -12,3 +20,92 @@ func TestServerHelp(t *testing.T) {
 		t.Fatalf("Command 'server --help' not working, got: %s", output)
 	}
 }
+
+type validator struct {
+	err error
+}
+
+func (v *validator) Check(actual, expected interface{}) {
+	if v.err != nil {
+		return
+	}
+	if actual != expected {
+		v.err = fmt.Errorf("Actual: '%v' is not equal to '%v'", actual, expected)
+	}
+}
+func (v *validator) Error() error { return v.err }
+
+type commandTest struct {
+	args           []string
+	env            map[string]string
+	cfgFileContent string
+	validation     func() error
+}
+
+func TestServerConfig(t *testing.T) {
+	actualDataDir, err := ioutil.TempDir("", "")
+	failErr(t, err, "making data dir")
+	tests := []commandTest{
+		{
+			args: []string{"server", "--data-dir", actualDataDir},
+			env:  map[string]string{"PILOSA_DATA_DIR": "/tmp/myEnvDatadir"},
+			cfgFileContent: `
+data-dir = "/tmp/myFileDatadir"
+host = "localhost:0"
+`,
+			validation: func() error {
+				v := validator{}
+				v.Check(cmd.Serve.Config.DataDir, actualDataDir)
+				v.Check(cmd.Serve.Config.Host, "localhost:0")
+				return v.Error()
+			},
+		},
+	}
+	for i, test := range tests {
+		com := setupCommand(t, test.args, test.env, test.cfgFileContent)
+		wait := sync.Mutex{}
+		wait.Lock()
+		var execErr error
+		go func() {
+			execErr = com.Execute()
+			wait.Unlock()
+		}()
+		// Serve.Close automatically waits for Serve.Run() to finish starting
+		// the server.
+		err := cmd.Serve.Close()
+		failErr(t, err, "closing pilosa server command")
+		wait.Lock() // make sure com.Execute finishes
+		failErr(t, execErr, "executing command")
+
+		if err := test.validation(); err != nil {
+			t.Fatalf("Failed test %d due to: %v", i, err)
+		}
+	}
+}
+
+func setupCommand(t *testing.T, args []string, env map[string]string, cfgFileContent string) *cobra.Command {
+	// make config file
+	cfgFile, err := ioutil.TempFile("", "")
+	failErr(t, err, "making temp file")
+	_, err = cfgFile.WriteString(cfgFileContent)
+	failErr(t, err, "writing config to temp file")
+
+	// set up config file args/env
+	env["PILOSA_CONFIG"] = cfgFile.Name()
+	args = append(args[:1], append([]string{"--config=" + cfgFile.Name()}, args[1:]...)...)
+
+	// set up env
+	for name, val := range env {
+		err = os.Setenv(name, val)
+		failErr(t, err, fmt.Sprintf("setting environment variable '%s' to '%s'", name, val))
+	}
+
+	// make command and set args
+	rc := cmd.NewRootCommand(strings.NewReader(""), ioutil.Discard, ioutil.Discard)
+	rc.SetArgs(args)
+
+	err = cfgFile.Close()
+	failErr(t, err, "closing config file")
+
+	return rc
+}

server/server.go

@@ -38,6 +38,11 @@ type Command struct {
 	Stdin  io.Reader
 	Stdout io.Writer
 	Stderr io.Writer
+
+	// running will be closed once Command.Run is finished.
+	running chan struct{}
+	// Done will be closed when Command.Close() is called
+	Done chan struct{}
 }
 
 // NewMain returns a new instance of Main.
@@ -49,11 +54,15 @@ func NewCommand() *Command {
 		Stdin:  os.Stdin,
 		Stdout: os.Stdout,
 		Stderr: os.Stderr,
+
+		running: make(chan struct{}),
+		Done:    make(chan struct{}),
 	}
 }
 
 // Run executes the pilosa server.
 func (m *Command) Run(args ...string) error {
+	defer close(m.running)
 	prefix := "~" + string(filepath.Separator)
 	if strings.HasPrefix(m.Config.DataDir, prefix) {
 		HomeDir := os.Getenv("HOME")
@@ -88,5 +97,9 @@ func (m *Command) Run(args ...string) error {
 
 // Close shuts down the server.
 func (m *Command) Close() error {
-	return m.Server.Close()
+	// must be running before it can be closed
+	<-m.running
+	err := m.Server.Close()
+	close(m.Done)
+	return err
 }