diff --git a/cmd/root.go b/cmd/root.go index 99888392a..1c1e3f2a4 100644 --- a/cmd/root.go +++ b/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 } diff --git a/cmd/root_test.go b/cmd/root_test.go index 3dec5908b..113815e45 100644 --- a/cmd/root_test.go +++ b/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) diff --git a/cmd/server.go b/cmd/server.go index 58497744e..f22404d9c 100644 --- a/cmd/server.go +++ b/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 } diff --git a/cmd/server_test.go b/cmd/server_test.go index c2a9edea0..4addb447b 100644 --- a/cmd/server_test.go +++ b/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 +} diff --git a/server/server.go b/server/server.go index 15f2ad9b4..1feb6c41b 100644 --- a/server/server.go +++ b/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 }