diff --git a/manager/workload/log_test.go b/manager/workload/log_test.go index e24b2f9..0bf60b4 100644 --- a/manager/workload/log_test.go +++ b/manager/workload/log_test.go @@ -3,16 +3,23 @@ package workload import ( "bufio" "context" + "net" "net/http" "testing" "time" "github.com/prometheus/client_golang/prometheus/testutil" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" "github.com/projecteru2/agent/types" ) +const ( + streamTimeout = 10 * time.Second + streamPoll = time.Millisecond +) + func TestLogBroadcaster(t *testing.T) { manager := newMockWorkloadManager(t) @@ -35,49 +42,44 @@ func TestLogBroadcaster(t *testing.T) { manager.PullLog(logCtx, app, buf) } } - server := &http.Server{Addr: ":12310"} - defer func() { _ = server.Shutdown(context.Background()) }() - go func() { - mux := http.NewServeMux() - mux.HandleFunc("GET /log/{$}", handler) - server.Handler = mux - assert.Equal(t, server.ListenAndServe(), http.ErrServerClosed) + listener, err := net.Listen("tcp", "127.0.0.1:0") + require.NoError(t, err) + + mux := http.NewServeMux() + mux.HandleFunc("GET /log/{$}", handler) + server := &http.Server{Handler: mux} + served := make(chan error, 1) + go func() { served <- server.Serve(listener) }() + defer func() { + assert.NoError(t, server.Shutdown(t.Context())) + assert.Equal(t, http.ErrServerClosed, <-served) }() - ctx, cancel := context.WithTimeout(t.Context(), 7*time.Second) - defer cancel() - - go func() { - time.Sleep(3 * time.Second) - manager.logBroadcaster.broadcast(ctx, &types.Log{ - ID: "Rei", - Name: "nerv", - Type: "stdout", - EntryPoint: "eva0", - Data: "data0", - }) - manager.logBroadcaster.broadcast(ctx, &types.Log{ - ID: "Rei", - Name: "nerv", - Type: "stdout", - EntryPoint: "eva0", - Data: "data1", - }) - }() - - time.Sleep(time.Second) - - reqCtx, reqCancel := context.WithTimeout(ctx, 3*time.Second) + reqCtx, reqCancel := context.WithTimeout(t.Context(), streamTimeout) defer reqCancel() - req, err := http.NewRequestWithContext(reqCtx, "GET", "http://127.0.0.1:12310/log/?app=nerv", nil) + req, err := http.NewRequestWithContext(reqCtx, http.MethodGet, "http://"+listener.Addr().String()+"/log/?app=nerv", nil) assert.Nil(t, err) resp, err := http.DefaultClient.Do(req) assert.Nil(t, err) defer resp.Body.Close() + require.Eventually(t, func() bool { + return subscriberCount(manager.logBroadcaster, "nerv") == 1 + }, streamTimeout, streamPoll, "the log stream never subscribed") + + for _, data := range []string{"data0", "data1"} { + manager.logBroadcaster.broadcast(t.Context(), &types.Log{ + ID: "Rei", + Name: "nerv", + Type: "stdout", + EntryPoint: "eva0", + Data: data, + }) + } + reader := bufio.NewReader(resp.Body) for range 2 { line, err := reader.ReadBytes('\n') @@ -86,9 +88,11 @@ func TestLogBroadcaster(t *testing.T) { } logCancel() - time.Sleep(time.Second) + require.Eventually(t, func() bool { + return subscriberCount(manager.logBroadcaster, "nerv") == 0 + }, streamTimeout, streamPoll, "the canceled log stream never detached") - manager.logBroadcaster.broadcast(ctx, &types.Log{ + manager.logBroadcaster.broadcast(t.Context(), &types.Log{ ID: "Rei", Name: "nerv", Type: "stdout", @@ -136,3 +140,9 @@ func TestBroadcastDoesNotBlockOnAStalledSubscriber(t *testing.T) { } assert.Positive(t, stalled.dropped.Load()) } + +func subscriberCount(l *logBroadcaster, app string) int { + l.mu.RLock() + defer l.mu.RUnlock() + return len(l.subscribersMap[app]) +} diff --git a/manager/workload/manager_test.go b/manager/workload/manager_test.go index 76beb2c..4eb53fa 100644 --- a/manager/workload/manager_test.go +++ b/manager/workload/manager_test.go @@ -3,6 +3,7 @@ package workload import ( "context" "testing" + "testing/synctest" "time" "github.com/stretchr/testify/assert" @@ -13,19 +14,31 @@ import ( "github.com/projecteru2/agent/types" ) +const ( + runTimeout = 30 * time.Second + connectTimeout = 5 * time.Second + journalDrainWait = 2 * connectTimeout +) + func TestRun(t *testing.T) { - manager := newMockWorkloadManager(t) - src := manager.source.(*mocks.Nerv) - ctx, cancel := context.WithTimeout(t.Context(), time.Second*30) - defer cancel() - go func() { - src.StartEvents() - src.StartCustomEvent(&types.WorkloadEventMessage{ - ID: "Kaworu", - Action: "start", - }) - }() - assert.Nil(t, manager.Run(ctx)) + t.Setenv("PATH", "") + synctest.Test(t, func(t *testing.T) { + manager := newMockWorkloadManager(t) + src := manager.source.(*mocks.Nerv) + ctx, cancel := context.WithTimeout(t.Context(), runTimeout) + defer cancel() + + go func() { + src.StartEvents() + src.StartCustomEvent(&types.WorkloadEventMessage{ + ID: "Kaworu", + Action: "start", + }) + }() + + assert.Nil(t, manager.Run(ctx)) + synctest.Sleep(journalDrainWait) + }) } func newMockWorkloadManager(t *testing.T) *Manager { @@ -44,7 +57,7 @@ func newMockWorkloadManager(t *testing.T) *Manager { Timeout: 5, CacheTTL: 300, }, - GlobalConnectionTimeout: 5 * time.Second, + GlobalConnectionTimeout: connectTimeout, } clients, err := manager.NewClients(t.Context(), config)