diff --git a/api/websocket_test.go b/api/websocket_test.go index d6c27ca..4f83b4d 100644 --- a/api/websocket_test.go +++ b/api/websocket_test.go @@ -310,20 +310,20 @@ func TestRawToCodeOceanWriter(t *testing.T) { func TestCodeOceanToRawReaderReturnsOnlyAfterOneByteWasRead(t *testing.T) { reader := newCodeOceanToRawReader(nil) - read := make(chan int) + read := make(chan bool) go func() { p := make([]byte, 10) - n, _ := reader.Read(p) - read <- n + _, _ = reader.Read(p) + read <- true }() t.Run("Does not return immediately when there is no data", func(t *testing.T) { - assert.False(t, waitForChannel(read, tests.ShortTimeout)) + assert.False(t, tests.ChannelReceivesSomething(read, tests.ShortTimeout)) }) t.Run("Returns when there is data available", func(t *testing.T) { reader.buffer <- byte(42) - assert.True(t, waitForChannel(read, tests.ShortTimeout)) + assert.True(t, tests.ChannelReceivesSomething(read, tests.ShortTimeout)) }) } @@ -343,34 +343,25 @@ func TestCodeOceanToRawReaderReturnsOnlyAfterOneByteWasReadFromConnection(t *tes cancel := reader.readInputLoop() defer cancel() - read := make(chan int) + read := make(chan bool) message := make([]byte, 10) go func() { - n, _ := reader.Read(message) - read <- n + _, _ = reader.Read(message) + read <- true }() t.Run("Does not return immediately when there is no data", func(t *testing.T) { - assert.False(t, waitForChannel(read, tests.ShortTimeout)) + assert.False(t, tests.ChannelReceivesSomething(read, tests.ShortTimeout)) }) t.Run("Returns when there is data available", func(t *testing.T) { messages <- strings.NewReader("Hello") - assert.True(t, waitForChannel(read, tests.ShortTimeout)) + assert.True(t, tests.ChannelReceivesSomething(read, tests.ShortTimeout)) }) } // --- Test suite specific test helpers --- -func waitForChannel(ch chan int, duration time.Duration) bool { - select { - case <-ch: - return true - case <-time.After(duration): - return false - } -} - func newNomadAllocationWithMockedApiClient(runnerId string) (r runner.Runner, mock *nomad.ExecutorAPIMock) { mock = &nomad.ExecutorAPIMock{} r = runner.NewNomadJob(runnerId, mock) diff --git a/nomad/nomad_test.go b/nomad/nomad_test.go index 49b0c3c..bf9c30e 100644 --- a/nomad/nomad_test.go +++ b/nomad/nomad_test.go @@ -759,10 +759,5 @@ func TestNullReaderDoesNotReturnImmediately(t *testing.T) { _, _ = reader.Read(p) close(readerReturned) }() - - select { - case <-readerReturned: - t.Fatal("Read should not return immediately.") - case <-time.After(tests.ShortTimeout): - } + assert.False(t, tests.ChannelReceivesSomething(readerReturned, tests.ShortTimeout)) } diff --git a/tests/util.go b/tests/util.go new file mode 100644 index 0000000..f72b4e4 --- /dev/null +++ b/tests/util.go @@ -0,0 +1,14 @@ +package tests + +import "time" + +// ChannelReceivesSomething waits timeout seconds for something to be received from channel ch. +// If something is received, it returns true. If the timeout expires without receiving anything, it return false. +func ChannelReceivesSomething(ch chan bool, timeout time.Duration) bool { + select { + case <-ch: + return true + case <-time.After(timeout): + return false + } +}