add tests and comments

Author: dylanhitt

command.go

@@ -145,7 +145,12 @@ func WithEnvironmentVariables(env EnvVars) func(c *Command) {
 	}
 }
 
+// WithContext adds context.Context to the command. Context timeout/deadline
+// is favored over Timeout
 func WithContext(ctx context.Context) func(c *Command) {
+	if ctx == nil {
+		panic("nil context")
+	}
 	return func(c *Command) {
 		c.ctx = ctx
 	}
@@ -219,10 +224,7 @@ func (c *Command) Execute() error {
 
 	done := make(chan error, 1)
 	go func() {
-		select {
-		case done <- cmd.Wait():
-			return
-		}
+		done <- cmd.Wait()
 	}()
 
 	select {
@@ -239,7 +241,6 @@ func (c *Command) Execute() error {
 			c.getExitCode(err)
 		}
 	}
-
 	c.executed = true
 	return nil
 }

command_test.go

@@ -2,12 +2,14 @@ package cmd
 
 import (
 	"bytes"
+	"context"
 	"fmt"
-	"github.com/stretchr/testify/assert"
 	"os"
 	"runtime"
 	"testing"
 	"time"
+
+	"github.com/stretchr/testify/assert"
 )
 
 func TestCommand_NewCommand(t *testing.T) {
@@ -129,3 +131,14 @@ func TestCommand_SetOptions(t *testing.T) {
 	assert.Equal(t, time.Duration(1000000000), c.Timeout)
 	assertEqualWithLineBreak(t, "test", writer.String())
 }
+
+func TestCommand_WithContext(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
+	defer cancel()
+	// WithContext is favored over Timeout
+	cmd := NewCommand("sleep 3;", WithContext(ctx), WithTimeout(1*time.Second))
+
+	err := cmd.Execute()
+	assert.NotNil(t, err)
+	assert.Equal(t, "context deadline exceeded", err.Error())
+}