Adding tests

dannykopping · dannykopping · commit 686c2b73b516 · 2024-03-15T14:14:19.000+02:00
Signed-off-by: Danny Kopping &lt;danny@coder.com&gt;
diff --git a/cli/cliui/provisionerjob_test.go b/cli/cliui/provisionerjob_test.go
@@ -2,8 +2,10 @@ package cliui_test
 
 import (
 	"context"
+	"fmt"
 	"io"
 	"os"
+	"regexp"
 	"runtime"
 	"sync"
 	"testing"
@@ -79,6 +81,70 @@ func TestProvisionerJob(t *testing.T) {
 		test.PTY.ExpectMatch("Something")
 	})
 
+	t.Run("Queue Position", func(t *testing.T) {
+		t.Parallel()
+
+		stage := cliui.ProvisioningStateQueued
+
+		tests := []struct {
+			name     string
+			queuePos int
+			expected string
+		}{
+			{
+				name:     "first",
+				queuePos: 0,
+				expected: fmt.Sprintf("%s$", stage),
+			},
+			{
+				name:     "next",
+				queuePos: 1,
+				expected: fmt.Sprintf(`%s %s$`, stage, regexp.QuoteMeta("(next)")),
+			},
+			{
+				name:     "other",
+				queuePos: 4,
+				expected: fmt.Sprintf(`%s %s$`, stage, regexp.QuoteMeta("(position: 4)")),
+			},
+		}
+
+		for _, tc := range tests {
+			tc := tc
+
+			t.Run(tc.name, func(t *testing.T) {
+				t.Parallel()
+
+				test := newProvisionerJob(t)
+				test.JobMutex.Lock()
+				test.Job.QueuePosition = tc.queuePos
+				test.Job.QueueSize = tc.queuePos
+				test.JobMutex.Unlock()
+
+				go func() {
+					<-test.Next
+					test.JobMutex.Lock()
+					test.Job.Status = codersdk.ProvisionerJobRunning
+					now := dbtime.Now()
+					test.Job.StartedAt = &now
+					test.JobMutex.Unlock()
+					<-test.Next
+					test.JobMutex.Lock()
+					test.Job.Status = codersdk.ProvisionerJobSucceeded
+					now = dbtime.Now()
+					test.Job.CompletedAt = &now
+					close(test.Logs)
+					test.JobMutex.Unlock()
+				}()
+				test.PTY.ExpectRegexMatch(tc.expected)
+				test.Next <- struct{}{}
+				test.PTY.ExpectMatch(cliui.ProvisioningStateQueued) // step completed
+				test.PTY.ExpectMatch(cliui.ProvisioningStateRunning)
+				test.Next <- struct{}{}
+				test.PTY.ExpectMatch(cliui.ProvisioningStateRunning)
+			})
+		}
+	})
+
 	// This cannot be ran in parallel because it uses a signal.
 	// nolint:paralleltest
 	t.Run("Cancel", func(t *testing.T) {
diff --git a/pty/ptytest/ptytest.go b/pty/ptytest/ptytest.go
@@ -6,6 +6,7 @@ import (
 	"context"
 	"fmt"
 	"io"
+	"regexp"
 	"runtime"
 	"strings"
 	"sync"
@@ -145,16 +146,36 @@ type outExpecter struct {
 }
 
 func (e *outExpecter) ExpectMatch(str string) string {
+	return e.expectMatchContext(str, e.ExpectMatchContext)
+}
+
+func (e *outExpecter) ExpectRegexMatch(str string) string {
+	return e.expectMatchContext(str, e.ExpectRegexMatchContext)
+}
+
+func (e *outExpecter) expectMatchContext(str string, impl func(ctx context.Context, str string) string) string {
 	e.t.Helper()
 
 	timeout, cancel := context.WithTimeout(context.Background(), testutil.WaitMedium)
 	defer cancel()
 
-	return e.ExpectMatchContext(timeout, str)
+	return impl(timeout, str)
 }
 
 // TODO(mafredri): Rename this to ExpectMatch when refactoring.
 func (e *outExpecter) ExpectMatchContext(ctx context.Context, str string) string {
+	return e.expectMatcher(ctx, str, func(src, pattern string) bool {
+		return strings.Contains(src, pattern)
+	})
+}
+
+func (e *outExpecter) ExpectRegexMatchContext(ctx context.Context, str string) string {
+	return e.expectMatcher(ctx, str, func(src, pattern string) bool {
+		return regexp.MustCompile(pattern).MatchString(src)
+	})
+}
+
+func (e *outExpecter) expectMatcher(ctx context.Context, str string, matchFn func(src, pattern string) bool) string {
 	e.t.Helper()
 
 	var buffer bytes.Buffer
@@ -168,7 +189,7 @@ func (e *outExpecter) ExpectMatchContext(ctx context.Context, str string) string
 			if err != nil {
 				return err
 			}
-			if strings.Contains(buffer.String(), str) {
+			if matchFn(buffer.String(), str) {
 				return nil
 			}
 		}
