github.com/polarismesh/polaris@v1.17.8/apiserver/grpcserver/base_test.go (about)

     1  /**
     2   * Tencent is pleased to support the open source community by making Polaris available.
     3   *
     4   * Copyright (C) 2019 THL A29 Limited, a Tencent company. All rights reserved.
     5   *
     6   * Licensed under the BSD 3-Clause License (the "License");
     7   * you may not use this file except in compliance with the License.
     8   * You may obtain a copy of the License at
     9   *
    10   * https://opensource.org/licenses/BSD-3-Clause
    11   *
    12   * Unless required by applicable law or agreed to in writing, software distributed
    13   * under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
    14   * CONDITIONS OF ANY KIND, either express or implied. See the License for the
    15   * specific language governing permissions and limitations under the License.
    16   */
    17  
    18  package grpcserver
    19  
    20  import (
    21  	"context"
    22  	"reflect"
    23  	"testing"
    24  
    25  	"google.golang.org/grpc/metadata"
    26  
    27  	"github.com/polarismesh/polaris/apiserver"
    28  	grpchelp "github.com/polarismesh/polaris/apiserver/grpcserver/utils"
    29  	"github.com/polarismesh/polaris/common/utils"
    30  )
    31  
    32  func mockGrpcContext(testVal map[string]string) context.Context {
    33  
    34  	md := make(metadata.MD)
    35  
    36  	for k := range testVal {
    37  		md[k] = []string{testVal[k]}
    38  	}
    39  
    40  	ctx := metadata.NewIncomingContext(context.Background(), md)
    41  
    42  	return ctx
    43  }
    44  
    45  func TestConvertContext(t *testing.T) {
    46  	type args struct {
    47  		ctx             context.Context
    48  		externalHeaders []string
    49  	}
    50  	tests := []struct {
    51  		name string
    52  		args args
    53  		want context.Context
    54  	}{
    55  		{
    56  			name: "",
    57  			args: args{
    58  				ctx: mockGrpcContext(map[string]string{
    59  					"internal-key-1": "internal-value-1",
    60  				}),
    61  			},
    62  			want: func() context.Context {
    63  				ctx := context.Background()
    64  
    65  				md := make(metadata.MD)
    66  
    67  				testVal := map[string]string{
    68  					"internal-key-1": "internal-value-1",
    69  				}
    70  
    71  				for k := range testVal {
    72  					md[k] = []string{testVal[k]}
    73  				}
    74  
    75  				ctx = context.WithValue(ctx, utils.ContextGrpcHeader, md)
    76  				ctx = context.WithValue(ctx, utils.StringContext("request-id"), "")
    77  				ctx = context.WithValue(ctx, utils.StringContext("client-ip"), "")
    78  				ctx = context.WithValue(ctx, utils.ContextClientAddress, "")
    79  				ctx = context.WithValue(ctx, utils.StringContext("user-agent"), "")
    80  
    81  				return ctx
    82  			}(),
    83  		},
    84  		{
    85  			name: "",
    86  			args: args{
    87  				ctx: mockGrpcContext(map[string]string{
    88  					"internal-key-1": "internal-value-1",
    89  					"request-id":     "request-id",
    90  					"user-agent":     "user-agent",
    91  				}),
    92  			},
    93  			want: func() context.Context {
    94  
    95  				md := make(metadata.MD)
    96  
    97  				testVal := map[string]string{
    98  					"internal-key-1": "internal-value-1",
    99  					"request-id":     "request-id",
   100  					"user-agent":     "user-agent",
   101  				}
   102  				for k := range testVal {
   103  					md[k] = []string{testVal[k]}
   104  				}
   105  
   106  				ctx := context.Background()
   107  				ctx = context.WithValue(ctx, utils.ContextGrpcHeader, md)
   108  				ctx = context.WithValue(ctx, utils.StringContext("request-id"), "request-id")
   109  				ctx = context.WithValue(ctx, utils.StringContext("client-ip"), "")
   110  				ctx = context.WithValue(ctx, utils.ContextClientAddress, "")
   111  				ctx = context.WithValue(ctx, utils.StringContext("user-agent"), "user-agent")
   112  
   113  				return ctx
   114  			}(),
   115  		},
   116  	}
   117  	for _, tt := range tests {
   118  		t.Run(tt.name, func(t *testing.T) {
   119  			if got := ConvertContext(tt.args.ctx); !reflect.DeepEqual(got, tt.want) {
   120  				t.Errorf("ConvertContext() = %v, \n want %v", got, tt.want)
   121  			}
   122  		})
   123  	}
   124  }
   125  
   126  func TestGetClientOpenMethod(t *testing.T) {
   127  	type args struct {
   128  		include  []string
   129  		protocol string
   130  	}
   131  	tests := []struct {
   132  		name    string
   133  		args    args
   134  		want    map[string]bool
   135  		wantErr bool
   136  	}{
   137  		{
   138  			name: "case=1",
   139  			args: args{
   140  				include: []string{
   141  					apiserver.RegisterAccess,
   142  				},
   143  				protocol: "grpc",
   144  			},
   145  			want: map[string]bool{
   146  				"/v1.PolarisGRPC/RegisterInstance":   true,
   147  				"/v1.PolarisGRPC/DeregisterInstance": true,
   148  			},
   149  			wantErr: false,
   150  		},
   151  		{
   152  			name: "case=1",
   153  			args: args{
   154  				include: []string{
   155  					apiserver.DiscoverAccess,
   156  				},
   157  				protocol: "grpc",
   158  			},
   159  			want: map[string]bool{
   160  				"/v1.PolarisGRPC/Discover":     true,
   161  				"/v1.PolarisGRPC/ReportClient": true,
   162  			},
   163  			wantErr: false,
   164  		},
   165  	}
   166  	for _, tt := range tests {
   167  		t.Run(tt.name, func(t *testing.T) {
   168  			got, err := grpchelp.GetClientOpenMethod(tt.args.include, tt.args.protocol)
   169  			if (err != nil) != tt.wantErr {
   170  				t.Errorf("GetClientOpenMethod() error = %v, wantErr %v", err, tt.wantErr)
   171  				return
   172  			}
   173  			if !reflect.DeepEqual(got, tt.want) {
   174  				t.Errorf("GetClientOpenMethod() = %v, want %v", got, tt.want)
   175  			}
   176  		})
   177  	}
   178  }