grpc_test.go

Documentation: go.opencensus.io/plugin/ocgrpc

     1  // Copyright 2017, OpenCensus Authors
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package ocgrpc
    16  
    17  import (
    18  	"context"
    19  	"sync"
    20  	"testing"
    21  	"time"
    22  
    23  	"google.golang.org/grpc/metadata"
    24  	"google.golang.org/grpc/stats"
    25  
    26  	"go.opencensus.io/stats/view"
    27  	"go.opencensus.io/trace"
    28  )
    29  
    30  func TestClientHandler(t *testing.T) {
    31  	ctx := context.Background()
    32  	te := &traceExporter{}
    33  	trace.RegisterExporter(te)
    34  	if err := view.Register(ClientSentMessagesPerRPCView); err != nil {
    35  		t.Fatal(err)
    36  	}
    37  	defer view.Unregister(ClientSentMessagesPerRPCView)
    38  
    39  	ctx, _ = trace.StartSpan(ctx, "/foo", trace.WithSampler(trace.AlwaysSample()))
    40  
    41  	var handler ClientHandler
    42  	ctx = handler.TagRPC(ctx, &stats.RPCTagInfo{
    43  		FullMethodName: "/service.foo/method",
    44  	})
    45  	handler.HandleRPC(ctx, &stats.Begin{
    46  		Client:    true,
    47  		BeginTime: time.Now(),
    48  	})
    49  	handler.HandleRPC(ctx, &stats.End{
    50  		Client:  true,
    51  		EndTime: time.Now(),
    52  	})
    53  
    54  	stats, err := view.RetrieveData(ClientSentMessagesPerRPCView.Name)
    55  	if err != nil {
    56  		t.Fatal(err)
    57  	}
    58  	traces := te.buffer
    59  
    60  	if got, want := len(stats), 1; got != want {
    61  		t.Errorf("Got %v stats; want %v", got, want)
    62  	}
    63  	if got, want := len(traces), 1; got != want {
    64  		t.Errorf("Got %v traces; want %v", got, want)
    65  	}
    66  }
    67  
    68  func TestServerHandler(t *testing.T) {
    69  	tests := []struct {
    70  		name         string
    71  		newTrace     bool
    72  		expectTraces int
    73  	}{
    74  		{"trust_metadata", false, 1},
    75  		{"no_trust_metadata", true, 0},
    76  	}
    77  
    78  	for _, test := range tests {
    79  		t.Run(test.name, func(t *testing.T) {
    80  
    81  			ctx := context.Background()
    82  
    83  			handler := &ServerHandler{
    84  				IsPublicEndpoint: test.newTrace,
    85  				StartOptions: trace.StartOptions{
    86  					Sampler: trace.ProbabilitySampler(0.0),
    87  				},
    88  			}
    89  
    90  			te := &traceExporter{}
    91  			trace.RegisterExporter(te)
    92  			if err := view.Register(ServerCompletedRPCsView); err != nil {
    93  				t.Fatal(err)
    94  			}
    95  
    96  			md := metadata.MD{
    97  				"grpc-trace-bin": []string{string([]byte{0, 0, 62, 116, 14, 118, 117, 157, 126, 7, 114, 152, 102, 125, 235, 34, 114, 238, 1, 187, 201, 24, 210, 231, 20, 175, 241, 2, 1})},
    98  			}
    99  			ctx = metadata.NewIncomingContext(ctx, md)
   100  			ctx = handler.TagRPC(ctx, &stats.RPCTagInfo{
   101  				FullMethodName: "/service.foo/method",
   102  			})
   103  			handler.HandleRPC(ctx, &stats.Begin{
   104  				BeginTime: time.Now(),
   105  			})
   106  			handler.HandleRPC(ctx, &stats.End{
   107  				EndTime: time.Now(),
   108  			})
   109  
   110  			rows, err := view.RetrieveData(ServerCompletedRPCsView.Name)
   111  			if err != nil {
   112  				t.Fatal(err)
   113  			}
   114  			traces := te.buffer
   115  
   116  			if got, want := len(rows), 1; got != want {
   117  				t.Errorf("Got %v rows; want %v", got, want)
   118  			}
   119  			if got, want := len(traces), test.expectTraces; got != want {
   120  				t.Errorf("Got %v traces; want %v", got, want)
   121  			}
   122  
   123  			// Cleanup.
   124  			view.Unregister(ServerCompletedRPCsView)
   125  		})
   126  	}
   127  }
   128  
   129  type traceExporter struct {
   130  	mu     sync.Mutex
   131  	buffer []*trace.SpanData
   132  }
   133  
   134  func (e *traceExporter) ExportSpan(sd *trace.SpanData) {
   135  	e.mu.Lock()
   136  	e.buffer = append(e.buffer, sd)
   137  	e.mu.Unlock()
   138  }
   139
View as plain text