diff --git a/flow/client/client.go b/flow/client/client.go index d74e4fd47..60057d746 100644 --- a/flow/client/client.go +++ b/flow/client/client.go @@ -86,6 +86,7 @@ func (c *GRPCClient) Close() error { if err := conn.Close(); err != nil && !errors.Is(err, context.Canceled) { return fmt.Errorf("close client connection: %w", err) } + return nil } @@ -108,7 +109,7 @@ func (c *GRPCClient) Send(event *proto.FlowEvent) error { func (c *GRPCClient) Receive(ctx context.Context, interval time.Duration, msgHandler func(msg *proto.FlowEventAck) error) error { backOff := defaultBackoff(ctx, interval) operation := func() error { - if err := c.establishStreamAndReceive(ctx, msgHandler); err == nil { + if err := c.establishStreamAndReceive(ctx, msgHandler); err != nil { if errors.Is(err, context.Canceled) { return backoff.Permanent(err) } diff --git a/flow/client/client_test.go b/flow/client/client_test.go index 0972c7b65..fcf08cb39 100644 --- a/flow/client/client_test.go +++ b/flow/client/client_test.go @@ -169,6 +169,11 @@ func TestReceive_ContextCancellation(t *testing.T) { assert.NoError(t, err, "failed to close flow") }) + go func() { + time.Sleep(100 * time.Millisecond) + cancel() + }() + handlerCalled := false msgHandler := func(msg *proto.FlowEventAck) error { if !msg.IsInitiator { @@ -276,3 +281,41 @@ func TestNewClient_PermanentClose(t *testing.T) { t.Fatal("Receive did not return after Close — stuck in retry loop") } } + +func TestNewClient_CloseVerify(t *testing.T) { + server := newTestServer(t) + + client, err := flow.NewClient("http://"+server.addr, "test-payload", "test-signature", 1*time.Second) + require.NoError(t, err) + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + t.Cleanup(cancel) + + done := make(chan error, 1) + go func() { + done <- client.Receive(ctx, 1*time.Second, func(msg *proto.FlowEventAck) error { + return nil + }) + }() + + closeDone := make(chan struct{}, 1) + go func() { + _ = client.Close() + closeDone <- struct{}{} + }() + + select { + case err := <-done: + require.Error(t, err) + case <-time.After(2 * time.Second): + t.Fatal("Receive did not return after Close — stuck in retry loop") + } + + select { + case <-closeDone: + return + case <-time.After(2 * time.Second): + t.Fatal("Receive did not return after Close — stuck in retry loop") + } + +}