• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1 /*
2  * Copyright 2022 The gRPC Authors
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *     http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 
17 package io.grpc.gcp.observability;
18 
19 import static com.google.common.truth.Truth.assertThat;
20 import static org.mockito.Mockito.mock;
21 import static org.mockito.Mockito.when;
22 
23 import com.google.cloud.trace.v1.TraceServiceClient;
24 import com.google.cloud.trace.v1.TraceServiceClient.ListTracesPagedResponse;
25 import com.google.devtools.cloudtrace.v1.GetTraceRequest;
26 import com.google.devtools.cloudtrace.v1.ListTracesRequest;
27 import com.google.devtools.cloudtrace.v1.Trace;
28 import com.google.devtools.cloudtrace.v1.TraceSpan;
29 import com.google.protobuf.util.Timestamps;
30 import io.grpc.ManagedChannelBuilder;
31 import io.grpc.Server;
32 import io.grpc.ServerBuilder;
33 import io.grpc.StaticTestingClassLoader;
34 import io.grpc.gcp.observability.interceptors.InternalLoggingChannelInterceptor;
35 import io.grpc.gcp.observability.interceptors.InternalLoggingServerInterceptor;
36 import io.grpc.gcp.observability.logging.GcpLogSink;
37 import io.grpc.gcp.observability.logging.Sink;
38 import io.grpc.testing.GrpcCleanupRule;
39 import io.grpc.testing.protobuf.SimpleServiceGrpc;
40 import io.opencensus.trace.samplers.Samplers;
41 import java.io.IOException;
42 import java.util.ArrayList;
43 import java.util.Collections;
44 import java.util.List;
45 import java.util.Map;
46 import java.util.concurrent.TimeUnit;
47 import java.util.regex.Pattern;
48 import org.junit.ClassRule;
49 import org.junit.Ignore;
50 import org.junit.Test;
51 import org.junit.runner.RunWith;
52 import org.junit.runners.JUnit4;
53 
54 @RunWith(JUnit4.class)
55 public class TracesTest {
56 
57   @ClassRule
58   public static final GrpcCleanupRule cleanupRule = new GrpcCleanupRule();
59 
60   private static final String PROJECT_ID = "PROJECT";
61   private static final String CUSTOM_TAG_KEY = "service";
62   private static final String CUSTOM_TAG_VALUE =
63       String.format("payment-%s", String.valueOf(System.currentTimeMillis()));
64   private static final Map<String, String> CUSTOM_TAGS =
65       Collections.singletonMap(CUSTOM_TAG_KEY, CUSTOM_TAG_VALUE);
66 
67   private final StaticTestingClassLoader classLoader =
68       new StaticTestingClassLoader(getClass().getClassLoader(),
69           Pattern.compile("io\\.grpc\\..*|io\\.opencensus\\..*"));
70 
71   /**
72    * End to end cloud trace test.
73    *
74    * <p>Ignoring test, because it calls external Cloud Tracing APIs. To test cloud trace setup
75    * locally,
76    * 1. Set up Cloud auth credentials
77    * 2. Assign permissions to service account to write traces to project specified by variable
78    * PROJECT_ID
79    * 3. Comment @Ignore annotation
80    * 4. This test is expected to pass when ran with above setup. This has been verified manually.
81    */
82   @Ignore
83   @Test
testTracesExporter()84   public void testTracesExporter() throws Exception {
85     Class<?> runnable =
86         classLoader.loadClass(TracesTest.StaticTestingClassTestTracesExporter.class.getName());
87     ((Runnable) runnable.getDeclaredConstructor().newInstance()).run();
88   }
89 
90   public static final class StaticTestingClassTestTracesExporter implements Runnable {
91 
92     @Override
run()93     public void run() {
94       Sink mockSink = mock(GcpLogSink.class);
95       ObservabilityConfig mockConfig = mock(ObservabilityConfig.class);
96       InternalLoggingChannelInterceptor.Factory mockChannelInterceptorFactory =
97           mock(InternalLoggingChannelInterceptor.Factory.class);
98       InternalLoggingServerInterceptor.Factory mockServerInterceptorFactory =
99           mock(InternalLoggingServerInterceptor.Factory.class);
100 
101       when(mockConfig.isEnableCloudTracing()).thenReturn(true);
102       when(mockConfig.getSampler()).thenReturn(Samplers.alwaysSample());
103       when(mockConfig.getProjectId()).thenReturn(PROJECT_ID);
104 
105       try {
106         GcpObservability observability =
107             GcpObservability.grpcInit(
108                 mockSink, mockConfig, mockChannelInterceptorFactory, mockServerInterceptorFactory);
109         observability.registerStackDriverExporter(PROJECT_ID, CUSTOM_TAGS);
110 
111         Server server =
112             ServerBuilder.forPort(0)
113                 .addService(new ObservabilityTestHelper.SimpleServiceImpl())
114                 .build()
115                 .start();
116         int port = cleanupRule.register(server).getPort();
117         SimpleServiceGrpc.SimpleServiceBlockingStub stub =
118             SimpleServiceGrpc.newBlockingStub(
119                 cleanupRule.register(
120                     ManagedChannelBuilder.forAddress("localhost", port).usePlaintext().build()));
121         assertThat(ObservabilityTestHelper.makeUnaryRpcViaClientStub("buddy", stub))
122             .isEqualTo("Hello buddy");
123         // Adding sleep to ensure traces are exported before querying cloud tracing backend
124         TimeUnit.SECONDS.sleep(10);
125 
126         TraceServiceClient traceServiceClient = TraceServiceClient.create();
127         String traceFilter =
128             String.format(
129                 "span:Sent.grpc.testing.SimpleService +%s:%s", CUSTOM_TAG_KEY, CUSTOM_TAG_VALUE);
130         String traceOrder = "start";
131         // Restrict time to last 1 minute
132         long startMillis = System.currentTimeMillis() - ((60 * 1) * 1000);
133         ListTracesRequest traceRequest =
134             ListTracesRequest.newBuilder()
135                 .setProjectId(PROJECT_ID)
136                 .setStartTime(Timestamps.fromMillis(startMillis))
137                 .setEndTime(Timestamps.fromMillis(System.currentTimeMillis()))
138                 .setFilter(traceFilter)
139                 .setOrderBy(traceOrder)
140                 .build();
141         ListTracesPagedResponse traceResponse = traceServiceClient.listTraces(traceRequest);
142         assertThat(traceResponse.iterateAll()).isNotEmpty();
143         List<String> traceIdList = new ArrayList<>();
144         for (Trace t : traceResponse.iterateAll()) {
145           traceIdList.add(t.getTraceId());
146         }
147 
148         for (String traceId : traceIdList) {
149           // This checks Cloud trace for the new trace that was just created.
150           GetTraceRequest getTraceRequest =
151               GetTraceRequest.newBuilder().setProjectId(PROJECT_ID).setTraceId(traceId).build();
152           Trace trace = traceServiceClient.getTrace(getTraceRequest);
153           assertThat(trace.getSpansList()).hasSize(3);
154           for (TraceSpan span : trace.getSpansList()) {
155             assertThat(span.getName()).contains("grpc.testing.SimpleService.UnaryRpc");
156             assertThat(span.getLabelsMap().get(CUSTOM_TAG_KEY)).isEqualTo(CUSTOM_TAG_VALUE);
157           }
158         }
159         observability.close();
160       } catch (IOException | InterruptedException e) {
161         throw new AssertionError("Exception while testing traces", e);
162       }
163     }
164   }
165 }
166