Improve test coverage for streaming exceptions (#61)

* [COVERAGE] Improve test coverage for streaming exceptions
* [BUGFIX] Ensure response streams close channels on exception (even when blocking)

Author: csvance

src/Curl.jl

@@ -42,6 +42,7 @@ function write_callback(
                 GRPC_INTERNAL,
                 "Recieved $(n) bytes from curl but only handled $(handled_n_bytes_total)",
             )
+            !isnothing(req.response_c) && close(req.response_c)
             return typemax(Csize_t)
         end
 
@@ -351,13 +352,15 @@ function handle_write(req::gRPCRequest, buf::Vector{UInt8})::Tuple{Int64, Union{
                     GRPC_UNIMPLEMENTED,
                     "Response was compressed but compression is not currently supported.",
                 )
+                !isnothing(req.response_c) && close(req.response_c)
                 notify(req.ready)
                 return n, nothing
             elseif req.response_length > req.max_recieve_message_length
                 req.ex = gRPCServiceCallException(
                     GRPC_RESOURCE_EXHAUSTED,
                     "length-prefix longer than max_recieve_message_length: $(req.response_length) > $(req.max_recieve_message_length)",
                 )
+                !isnothing(req.response_c) && close(req.response_c)
                 notify(req.ready)
                 return n, nothing
             end

src/Streaming.jl

@@ -57,9 +57,6 @@ function grpc_async_stream_request(
             end
         end
     catch ex
-        close(channel)
-        close(req.request_c)
-
         if isa(ex, InvalidStateException)
             # Wait for any request data to be flushed by curl
             wait(req.curl_done_reading)
@@ -82,6 +79,9 @@ function grpc_async_stream_request(
             end
             @error "grpc_async_stream_request: unexpected exception" exception = ex
         end
+    finally
+        close(channel)
+        close(req.request_c)
     end
 
     nothing
@@ -98,9 +98,6 @@ function grpc_async_stream_response(
             put!(channel, response)
         end
     catch ex
-        close(channel)
-        close(req.response_c)
-
         if isa(ex, InvalidStateException)
 
         elseif isa(ex, gRPCServiceCallException)
@@ -115,7 +112,9 @@ function grpc_async_stream_response(
             end
             @error "grpc_async_stream_response: unexpected exception" exception = ex
         end
-
+    finally
+        close(channel)
+        close(req.response_c)
     end
 
     nothing

test/runtests.jl

@@ -239,18 +239,110 @@ include("gen/test/test_pb.jl")
         @test length(response.data) == N
     # end
 
-    # @testset "Don't Stick User Tasks" 
+    # @testset "Don't Stick User Tasks"
         # This fails on Julia 1.10 but works on Julia 1.12
         client = TestService_TestRPC_Client(_TEST_HOST, _TEST_PORT)
 
-        task = @sync begin 
-            @spawn begin 
+        task = @sync begin
+            @spawn begin
                 grpc_sync_request(client, TestRequest(1, zeros(UInt64, 1)))
             end
-        end 
+        end
 
         @test !task.sticky
     # end
+
+    # @testset "grpc_async_stream_request - gRPCServiceCallException" begin
+        # Test that gRPCServiceCallException is properly stored in req.ex
+        client = TestService_TestClientStreamRPC_Client(_TEST_HOST, _TEST_PORT; max_send_message_length=100)
+        request_c = Channel{TestRequest}(1)
+
+        req = grpc_async_request(client, request_c)
+
+        # Send a request that exceeds max_send_message_length to trigger gRPCServiceCallException
+        put!(request_c, TestRequest(1, zeros(UInt64, 1000)))
+        close(request_c)
+
+        # Wait and check that the exception is a gRPCServiceCallException
+        try
+            grpc_async_await(client, req)
+            @test false  # Should not reach here
+        catch ex
+            @test isa(ex, gRPCServiceCallException)
+        end
+    # end
+
+    # @testset "grpc_async_stream_request - general exception" begin
+        # Test the else branch with a non-gRPC exception
+        client = TestService_TestClientStreamRPC_Client(_TEST_HOST, _TEST_PORT)
+        request_c = Channel{TestRequest}(1)
+
+        req = grpc_async_request(client, request_c)
+
+        # Close the channel and then try to take from it (triggers InvalidStateException)
+        close(request_c)
+
+        # Give the async task time to encounter the exception
+        sleep(0.2)
+
+        # The InvalidStateException should be handled gracefully
+        # and the request should complete (possibly with no error or a different error)
+        try
+            grpc_async_await(client, req)
+        catch ex
+            # If there's an exception, it shouldn't be InvalidStateException
+            # (that should be handled internally)
+            @test !isa(ex, InvalidStateException)
+        end
+    # end
+
+    # @testset "grpc_async_stream_response - InvalidStateException" begin
+        # Test that InvalidStateException is handled when response channel closes early
+        client = TestService_TestServerStreamRPC_Client(_TEST_HOST, _TEST_PORT)
+        response_c = Channel{TestResponse}(1)
+
+        req = grpc_async_request(client, TestRequest(10, zeros(UInt64, 1)), response_c)
+
+        # Take one response then close the channel to trigger InvalidStateException in put!
+        response = take!(response_c)
+        @test length(response.data) >= 1
+        close(response_c)
+
+        # Give time for the async task to encounter InvalidStateException
+        sleep(0.2)
+
+        # InvalidStateException should be handled internally without propagating
+        try
+            grpc_async_await(req)
+        catch ex
+            # If there's an exception, it shouldn't be InvalidStateException
+            @test !isa(ex, InvalidStateException)
+        end
+    # end
+
+    # @testset "grpc_async_stream_response - gRPCServiceCallException" begin
+        # Test that gRPCServiceCallException is properly handled in response stream
+        # Use a client with restrictive max_recieve_message_length
+        client = TestService_TestServerStreamRPC_Client(_TEST_HOST, _TEST_PORT; max_recieve_message_length=1)
+        response_c = Channel{TestResponse}(100)
+
+        # Request a response that will exceed the max size
+        req = grpc_async_request(client, TestRequest(10, zeros(UInt64, 100)), response_c)
+
+        # Wait for the error to occur
+        sleep(0.2)
+
+        # Should get gRPCServiceCallException when awaiting
+        try
+            for response in response_c
+                # Might get some responses before the error
+            end
+            grpc_async_await(req)
+            @test false  # Should not reach here
+        catch ex
+            @test isa(ex, gRPCServiceCallException)
+        end
+    # end
     end
 
     grpc_shutdown()