Improve tests

Signed-off-by: Tim Vaillancourt <[email protected]>

go/logic/server.go

@@ -23,6 +23,8 @@ import (
 	"github.com/github/gh-ost/go/base"
 )
 
+const defaultCPUProfileDuration = time.Second * 30
+
 var ErrCPUProfilingInProgress = errors.New("cpu profiling already in progress")
 
 type printStatusFunc func(PrintStatusRule, io.Writer)
@@ -45,16 +47,24 @@ func NewServer(migrationContext *base.MigrationContext, hooksExecutor *HooksExec
 	}
 }
 
-func (this *Server) runCPUProfile(duration time.Duration) (string, error) {
+func (this *Server) runCPUProfile(arg string) (string, error) {
 	if atomic.LoadInt64(&this.isCPUProfiling) > 0 {
 		return "", ErrCPUProfilingInProgress
 	}
 
+	var err error
+	duration := defaultCPUProfileDuration
+	if arg != "" {
+		if duration, err = time.ParseDuration(arg); err != nil {
+			return "", err
+		}
+	}
+
 	atomic.StoreInt64(&this.isCPUProfiling, 1)
 	defer atomic.StoreInt64(&this.isCPUProfiling, 0)
 
 	var buf bytes.Buffer
-	if err := pprof.StartCPUProfile(&buf); err != nil {
+	if err = pprof.StartCPUProfile(&buf); err != nil {
 		return "", err
 	}
 
@@ -198,13 +208,7 @@ help                                 # This message
 	case "info", "status":
 		return ForcePrintStatusAndHintRule, nil
 	case "cpu-profile":
-		profileDuration := time.Second * 30
-		if arg != "" {
-			if profileDuration, err = time.ParseDuration(arg); err != nil {
-				return NoPrintStatusRule, err
-			}
-		}
-		profile, err := this.runCPUProfile(profileDuration)
+		profile, err := this.runCPUProfile(arg)
 		if err == nil {
 			fmt.Fprintln(writer, profile)
 		}

go/logic/server_test.go

@@ -3,15 +3,20 @@ package logic
 import (
 	"encoding/base64"
 	"testing"
-	"time"
 
 	"github.com/openark/golib/tests"
 )
 
 func TestRunCPUProfile(t *testing.T) {
 	{
 		s := &Server{isCPUProfiling: 0}
-		profile, err := s.runCPUProfile(time.Millisecond * 10)
+		profile, err := s.runCPUProfile("should-fail")
+		tests.S(t).ExpectNotNil(err)
+		tests.S(t).ExpectEquals(profile, "")
+	}
+	{
+		s := &Server{isCPUProfiling: 0}
+		profile, err := s.runCPUProfile("10ms")
 		tests.S(t).ExpectNil(err)
 		tests.S(t).ExpectNotEquals(profile, "")
 
@@ -21,8 +26,8 @@ func TestRunCPUProfile(t *testing.T) {
 	}
 	{
 		s := &Server{isCPUProfiling: 1}
-		profile, err := s.runCPUProfile(time.Millisecond * 10)
-		tests.S(t).ExpectNotNil(err)
+		profile, err := s.runCPUProfile("15ms")
+		tests.S(t).ExpectEquals(err, ErrCPUProfilingInProgress)
 		tests.S(t).ExpectEquals(profile, "")
 	}
 }