tortoise/main_test.go

package tortoise_test

import (
	"context"
	"os/exec"
	"sync"
	"testing"
	"time"

	"git.iamthefij.com/iamthefij/tortoise"
)

const (
	TaskStartWait = 10 * time.Millisecond
)

func TestShellRunnerNoCallback(t *testing.T) {
	t.Parallel()

	cases := []struct {
		command    string
		output     string
		ReturnCode int
	}{
		{"echo hello world", "hello world\n", 0},
		{"echo hello world && exit 1", "hello world\n", 1},
	}

	for _, c := range cases {
		c := c
		t.Run(c.command, func(t *testing.T) {
			t.Parallel()

			runner := tortoise.NewShellRunner()
			runner.Start()

			// Test command without callback
			if err := runner.AddCommand(c.command, nil); err != nil {
				t.Fatalf("unexpected error adding command: %v", err)
			}

			// Wait a sec for the worker to pick up the task
			time.Sleep(TaskStartWait)

			runner.Stop()

			result := runner.GetResults()
			if result == nil {
				t.Fatal("expected result, got nil")
			}

			if result.Output != c.output || result.ReturnCode != c.ReturnCode {
				t.Fatalf("expected output '%s' and return code %d, got '%s' and %d", c.output, c.ReturnCode, result.Output, result.ReturnCode)
			}
		})
	}
}

func TestShellRunnerCallback(t *testing.T) {
	t.Parallel()

	runner := tortoise.NewShellRunner()
	runner.Start()

	// Test command with callback
	outputString := ""
	callbackWait := sync.WaitGroup{}

	callbackWait.Add(1)

	if err := runner.AddCommand("echo callback a", func(result *tortoise.CommandResult) {
		if result.Output != "callback a\n" {
			t.Fatalf("expected 'callback a', got '%s'", result.Output)
		}
		outputString = outputString + "a"
		callbackWait.Done()
	}); err != nil {
		t.Fatalf("unexpected error adding command: %v", err)
	}

	callbackWait.Add(1)

	if err := runner.AddCommand("echo callback b", func(result *tortoise.CommandResult) {
		if result.Output != "callback b\n" {
			t.Fatalf("expected 'callback b', got '%s'", result.Output)
		}
		outputString = outputString + "b"
		callbackWait.Done()
	}); err != nil {
		t.Fatalf("unexpected error adding command: %v", err)
	}

	// Wait a sec for the worker to pick up the task
	time.Sleep(TaskStartWait)

	// Timeout waiting for callbacks
	done := make(chan struct{})
	go func() {
		callbackWait.Wait()
		close(done)
	}()

	select {
	case <-done:
	case <-time.After(2 * time.Second):
		t.Fatal("callbacks timed out")
	}

	if outputString != "ab" {
		t.Fatal("callbacks were not reached in order:", outputString)
	}

	runner.Stop()

	// Make sure stop and kill both exit gracefully after the runner is stopped
	runner.Stop()
	runner.Kill()
}

func TestShellRunnerKillWithTimeout(t *testing.T) {
	t.Parallel()

	runner := tortoise.NewShellRunner()
	runner.Start()

	// Test command with callback
	callbackReached := false

	if err := runner.AddCommand("sleep 10 && echo callback test", func(result *tortoise.CommandResult) {
		callbackReached = true
		if result.Output != "callback test\n" {
			t.Fatalf("expected 'callback test', got '%s'", result.Output)
		}
	}); err != nil {
		t.Fatalf("unexpected error adding command: %v", err)
	}

	// Wait a sec for the worker to pick up the task
	time.Sleep(TaskStartWait)

	if err := runner.KillWithTimeout(1 * time.Second); err == nil {
		t.Fatal("expected error when killing commands, but got none")
	}

	if callbackReached {
		t.Fatal("callback was reached before kill")
	}
}

func TestStopPreventsNewCommands(t *testing.T) {
	runner := tortoise.NewShellRunner()
	runner.Start()

	runner.Stop()

	err := runner.AddCommand("echo should not run", nil)
	if err == nil {
		t.Fatal("expected error when adding command after stop, but got none")
	}
}

func TestAddingPriorToStart(t *testing.T) {
	runner := tortoise.NewShellRunner()

	err := runner.AddCommand("echo should not run", nil)
	if err == nil {
		t.Fatal("Should have failed to add prior to starting runner")
	}
}

func TestAddCmdWithTimeout(t *testing.T) {
	t.Parallel()

	runner := tortoise.NewShellRunner()
	runner.Start()

	ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)

	cmd := exec.CommandContext(ctx, "sleep", "10")

	err := runner.AddCmd(cmd, nil, cancel)
	if err != nil {
		t.Fatalf("unexpected error adding command: %v", err)
	}

	// Wait a sec for the worker to pick up the task
	time.Sleep(TaskStartWait)

	runner.Stop()

	result := runner.GetResults()
	if result == nil {
		t.Fatal("expected result, got nil")
	}

	if result.ReturnCode != -1 {
		t.Fatalf("expected return code -1, got %d", result.ReturnCode)
	}
}

func TestShellWithTimeout(t *testing.T) {
	t.Parallel()

	runner := tortoise.NewShellRunner()
	runner.SetTimeout(1 * time.Second)
	runner.Start()

	err := runner.AddCommand("sleep 10", nil)
	if err != nil {
		t.Fatalf("unexpected error adding command: %v", err)
	}

	// Wait a sec for the worker to pick up the task
	time.Sleep(TaskStartWait)

	runner.Stop()

	result := runner.GetResults()
	if result == nil {
		t.Fatal("expected result, got nil")
	}

	if result.ReturnCode != -1 {
		t.Fatalf("expected return code -1, got %d", result.ReturnCode)
	}
}