Refactor praefect server tests

1 files changed, 10 insertions, 3 deletions

diff --git a/internal/praefect/auth_test.go b/internal/praefect/auth_test.go
index f7b27394b..25fe7e4a0 100644
--- a/internal/praefect/auth_test.go
+++ b/internal/praefect/auth_test.go
@@ -64,7 +64,7 @@ func TestAuthFailures(t *testing.T) {
 
 			cli := mock.NewSimpleServiceClient(conn)
 
-			_, err = cli.SimpleUnaryUnary(ctx, &mock.SimpleRequest{
+			_, err = cli.ServerAccessor(ctx, &mock.SimpleRequest{
 				Value: 1,
 			})
 
@@ -133,7 +133,7 @@ func TestAuthSuccess(t *testing.T) {
 
 			cli := mock.NewSimpleServiceClient(conn)
 
-			_, err = cli.SimpleUnaryUnary(ctx, &mock.SimpleRequest{
+			_, err = cli.ServerAccessor(ctx, &mock.SimpleRequest{
 				Value: 1,
 			})
 
@@ -155,7 +155,14 @@ func dial(serverSocketPath string, opts []grpc.DialOption) (*grpc.ClientConn, er
 
 func runServer(t *testing.T, token string, required bool) (*Server, string, func()) {
 	backendToken := "abcxyz"
-	backend, cleanup := newMockDownstream(t, backendToken, callbackIncrement)
+	mockServer := &mockSvc{
+		serverAccessor: func(_ context.Context, req *mock.SimpleRequest) (*mock.SimpleResponse, error) {
+			return &mock.SimpleResponse{
+				Value: req.Value + 1,
+			}, nil
+		},
+	}
+	backend, cleanup := newMockDownstream(t, backendToken, mockServer)
 
 	conf := config.Config{
 		VirtualStorageName: "praefect",