pkg/rpcserver: refactor to remove Fatalf calls

Apply necessary changes to pkg/flatrpc and pkg/manager as well.

1 files changed, 7 insertions, 1 deletions

diff --git a/pkg/rpcserver/rpcserver_test.go b/pkg/rpcserver/rpcserver_test.go
index a885ad720..0c4984e93 100644
--- a/pkg/rpcserver/rpcserver_test.go
+++ b/pkg/rpcserver/rpcserver_test.go
@@ -4,10 +4,12 @@
 package rpcserver
 
 import (
+	"context"
 	"net"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+	"golang.org/x/sync/errgroup"
 
 	"github.com/google/syzkaller/pkg/flatrpc"
 	"github.com/google/syzkaller/pkg/mgrconfig"
@@ -212,7 +214,11 @@ func TestHandleConn(t *testing.T) {
 			serv.CreateInstance(1, injectExec, nil)
 
 			go flatrpc.Send(clientConn, tt.req)
-			serv.handleConn(serverConn)
+			var eg errgroup.Group
+			serv.handleConn(context.Background(), &eg, serverConn)
+			if err := eg.Wait(); err != nil {
+				t.Fatal(err)
+			}
 		})
 	}
 }