1 files changed, 67 insertions, 0 deletions
diff --git a/internal/handler/exec_test.go b/internal/handler/exec_test.go
index ce672f4..0dbd538 100644
--- a/internal/handler/exec_test.go
+++ b/internal/handler/exec_test.go
@@ -9,7 +9,10 @@ import (
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/metadata"
 
+	pb "gitlab.com/gitlab-org/gitaly/proto/go/gitalypb"
 	"gitlab.com/gitlab-org/gitlab-shell/internal/config"
+	"gitlab.com/gitlab-org/gitlab-shell/internal/gitlabnet/accessverifier"
+	"gitlab.com/gitlab-org/gitlab-shell/internal/testhelper"
 )
 
 func makeHandler(t *testing.T, err error) func(context.Context, *grpc.ClientConn) (int32, error) {
@@ -83,3 +86,67 @@ func TestGetConnMetadata(t *testing.T) {
 		})
 	}
 }
+
+func TestPrepareContext(t *testing.T) {
+	tests := []struct {
+		name             string
+		gc               *GitalyCommand
+		sshConnectionEnv string
+		repo             *pb.Repository
+		response         *accessverifier.Response
+		want             map[string]string
+	}{
+		{
+			name: "client_identity",
+			gc: &GitalyCommand{
+				Config:  &config.Config{},
+				Address: "tcp://localhost:9999",
+			},
+			sshConnectionEnv: "10.0.0.1 1234 127.0.0.1 5678",
+			repo: &pb.Repository{
+				StorageName:                   "default",
+				RelativePath:                  "@hashed/5f/9c/5f9c4ab08cac7457e9111a30e4664920607ea2c115a1433d7be98e97e64244ca.git",
+				GitObjectDirectory:            "path/to/git_object_directory",
+				GitAlternateObjectDirectories: []string{"path/to/git_alternate_object_directory"},
+				GlRepository:                  "project-26",
+				GlProjectPath:                 "group/private",
+			},
+			response: &accessverifier.Response{
+				KeyId:    1,
+				KeyType:  "key",
+				UserId:   "6",
+				Username: "jane.doe",
+			},
+			want: map[string]string{
+				"key_id":    "1",
+				"key_type":  "key",
+				"user_id":   "6",
+				"username":  "jane.doe",
+				"remote_ip": "10.0.0.1",
+			},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			cleanup, err := testhelper.Setenv("SSH_CONNECTION", tt.sshConnectionEnv)
+			require.NoError(t, err)
+			defer cleanup()
+
+			ctx := context.Background()
+
+			ctx, cancel := tt.gc.PrepareContext(ctx, tt.repo, tt.response, "protocol")
+			defer cancel()
+
+			md, exists := metadata.FromOutgoingContext(ctx)
+			require.True(t, exists)
+			require.Equal(t, len(tt.want), md.Len())
+
+			for k, v := range tt.want {
+				values := md.Get(k)
+				require.Equal(t, 1, len(values))
+				require.Equal(t, v, values[0])
+			}
+
+		})
+	}
+}