add WaitFor test

app_test.go

@@ -303,3 +303,60 @@ func TestPIDRemoved(t *testing.T) {
 		})
 	}
 }
+
+func TestWaitFor(t *testing.T) {
+	tests := []struct {
+		name   string
+		ctx    func(t *testing.T) context.Context
+		ch     func(t *testing.T) chan int
+		exp    int
+		expErr error
+	}{
+		{
+			name: "canceled with io error",
+			ctx: func(t *testing.T) context.Context {
+				ctx, cancel := context.WithCancelCause(context.Background())
+				cancel(io.ErrUnexpectedEOF)
+				return ctx
+			},
+			ch: func(t *testing.T) chan int {
+				return nil
+			},
+			exp:    0,
+			expErr: io.ErrUnexpectedEOF,
+		},
+		{
+			name: "canceled",
+			ctx: func(t *testing.T) context.Context {
+				ctx, cancel := context.WithCancel(context.Background())
+				cancel()
+				return ctx
+			},
+			ch: func(t *testing.T) chan int {
+				return nil
+			},
+			exp:    0,
+			expErr: context.Canceled,
+		},
+		{
+			name: "success",
+			ctx: func(t *testing.T) context.Context {
+				return context.Background()
+			},
+			ch: func(t *testing.T) chan int {
+				ch := make(chan int, 1)
+				ch <- 1
+				close(ch)
+				return ch
+			},
+			exp: 1,
+		},
+	}
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			i, err := lu.WaitFor(tc.ctx(t), tc.ch(t))
+			jtest.Require(t, tc.expErr, err)
+			assert.Equal(t, tc.exp, i)
+		})
+	}
+}