33
33
import com .google .protobuf .ByteString ;
34
34
import com .google .protobuf .BytesValue ;
35
35
import com .google .protobuf .StringValue ;
36
+ import io .grpc .Metadata ;
36
37
import io .grpc .Server ;
37
38
import io .grpc .ServerBuilder ;
39
+ import io .grpc .ServerCall ;
40
+ import io .grpc .ServerCall .Listener ;
41
+ import io .grpc .ServerCallHandler ;
42
+ import io .grpc .ServerInterceptor ;
38
43
import io .grpc .stub .StreamObserver ;
39
44
import java .io .IOException ;
40
45
import java .net .ServerSocket ;
@@ -56,6 +61,7 @@ public class EnhancedBigtableStubTest {
56
61
private static final String APP_PROFILE_ID = "app-profile-id" ;
57
62
58
63
private Server server ;
64
+ private MetadataInterceptor metadataInterceptor ;
59
65
private FakeDataService fakeDataService ;
60
66
private EnhancedBigtableStubSettings defaultSettings ;
61
67
private EnhancedBigtableStub enhancedBigtableStub ;
@@ -66,8 +72,13 @@ public void setUp() throws IOException, IllegalAccessException, InstantiationExc
66
72
try (ServerSocket ss = new ServerSocket (0 )) {
67
73
port = ss .getLocalPort ();
68
74
}
75
+ metadataInterceptor = new MetadataInterceptor ();
69
76
fakeDataService = new FakeDataService ();
70
- server = ServerBuilder .forPort (port ).addService (fakeDataService ).build ();
77
+ server =
78
+ ServerBuilder .forPort (port )
79
+ .intercept (metadataInterceptor )
80
+ .addService (fakeDataService )
81
+ .build ();
71
82
server .start ();
72
83
73
84
defaultSettings =
@@ -137,6 +148,33 @@ public void testChannelPrimerConfigured() throws IOException {
137
148
}
138
149
}
139
150
151
+ @ Test
152
+ public void testUserAgent () throws InterruptedException {
153
+ ServerStreamingCallable <Query , Row > streamingCallable =
154
+ enhancedBigtableStub .createReadRowsCallable (new DefaultRowAdapter ());
155
+
156
+ Query request = Query .create ("table-id" ).rowKey ("row-key" );
157
+ streamingCallable .call (request ).iterator ().next ();
158
+
159
+ assertThat (metadataInterceptor .headers ).hasSize (1 );
160
+ Metadata metadata = metadataInterceptor .headers .take ();
161
+ assertThat (metadata .get (Metadata .Key .of ("user-agent" , Metadata .ASCII_STRING_MARSHALLER )))
162
+ .contains ("bigtable-java/" );
163
+ }
164
+
165
+ private static class MetadataInterceptor implements ServerInterceptor {
166
+ final BlockingQueue <Metadata > headers = Queues .newLinkedBlockingDeque ();
167
+
168
+ @ Override
169
+ public <ReqT , RespT > Listener <ReqT > interceptCall (
170
+ ServerCall <ReqT , RespT > serverCall ,
171
+ Metadata metadata ,
172
+ ServerCallHandler <ReqT , RespT > serverCallHandler ) {
173
+ headers .add (metadata );
174
+ return serverCallHandler .startCall (serverCall , metadata );
175
+ }
176
+ }
177
+
140
178
private static class FakeDataService extends BigtableGrpc .BigtableImplBase {
141
179
final BlockingQueue <ReadRowsRequest > requests = Queues .newLinkedBlockingDeque ();
142
180
0 commit comments