Fix test

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)
 			}

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")
+	}
+
+}