11package org .hypertrace .core .grpcutils .client ;
22
3+ import static org .mockito .Mockito .mock ;
4+
35import io .grpc .CallCredentials ;
46import io .grpc .Context ;
57import io .grpc .Metadata ;
1214import org .junit .jupiter .api .Assertions ;
1315import org .junit .jupiter .api .Test ;
1416
15- import static io .grpc .Metadata .BINARY_BYTE_MARSHALLER ;
16- import static org .mockito .Mockito .mock ;
17-
1817public class ContextKeyBasedCredsTest {
1918 private static final String TENANT_ID = "test-tenant-id" ;
2019 private static final String TEST_AUTH_HEADER = "test-auth-header" ;
2120
2221 @ Test
2322 public void testApplyRequestMetadata_shouldApplyAllHeaders () {
24- CallCredentials callCredentials = RequestContextClientCallCredsProviderFactory .getClientCallCredsProvider ().get ();
23+ CallCredentials callCredentials =
24+ RequestContextClientCallCredsProviderFactory .getClientCallCredsProvider ().get ();
2525
2626 Assertions .assertTrue (callCredentials instanceof ContextKeyBasedCreds );
2727
@@ -32,45 +32,57 @@ public void testApplyRequestMetadata_shouldApplyAllHeaders() {
3232 requestContext .add ("x-some-tenant-header" , "v1" );
3333 requestContext .add ("grpc-trace-bin" , "AAARf5ZpQwlN/8FVe1axOPlaAQIdRU/Y8j0LAgE" );
3434
35-
3635 Context ctx = Context .current ().withValue (RequestContext .CURRENT , requestContext );
3736
3837 Context previous = ctx .attach ();
3938 try {
40- callCredentials .applyRequestMetadata (mock (CallCredentials .RequestInfo .class ), mock (Executor .class ), new CallCredentials .MetadataApplier () {
41- @ Override
42- public void apply (Metadata headers ) {
43- Map <String , String > headersMap = headers .keys ().stream ()
44- .collect (Collectors .toUnmodifiableMap (k -> k , k -> {
45- String value ;
46- if (k .toLowerCase ().endsWith (Metadata .BINARY_HEADER_SUFFIX )) {
47- value = new String (headers .get (Metadata .Key .of (k , Metadata .BINARY_BYTE_MARSHALLER )));
48- } else {
49- value = headers .get (Metadata .Key .of (k , Metadata .ASCII_STRING_MARSHALLER ));
50- }
51- return value ;
52- } ));
39+ callCredentials .applyRequestMetadata (
40+ mock (CallCredentials .RequestInfo .class ),
41+ mock (Executor .class ),
42+ new CallCredentials .MetadataApplier () {
43+ @ Override
44+ public void apply (Metadata headers ) {
45+ Map <String , String > headersMap =
46+ headers .keys ().stream ()
47+ .collect (
48+ Collectors .toUnmodifiableMap (
49+ k -> k ,
50+ k -> {
51+ String value ;
52+ if (k .toLowerCase ().endsWith (Metadata .BINARY_HEADER_SUFFIX )) {
53+ value =
54+ new String (
55+ headers .get (
56+ Metadata .Key .of (k , Metadata .BINARY_BYTE_MARSHALLER )));
57+ } else {
58+ value =
59+ headers .get (
60+ Metadata .Key .of (k , Metadata .ASCII_STRING_MARSHALLER ));
61+ }
62+ return value ;
63+ }));
5364
54- // Should filter out the TENANT_ID_HEADER_KEY
55- Assertions .assertEquals (
56- Map .of (
57- RequestContextConstants .TENANT_ID_HEADER_KEY , TENANT_ID ,
58- RequestContextConstants .AUTHORIZATION_HEADER , TEST_AUTH_HEADER ,
59- "x-some-tenant-header" , "v1" ,
60- "grpc-trace-bin" , "AAARf5ZpQwlN/8FVe1axOPlaAQIdRU/Y8j0LAgE"
61- ),
62- headersMap );
63- }
65+ // Should filter out the TENANT_ID_HEADER_KEY
66+ Assertions .assertEquals (
67+ Map .of (
68+ RequestContextConstants .TENANT_ID_HEADER_KEY ,
69+ TENANT_ID ,
70+ RequestContextConstants .AUTHORIZATION_HEADER ,
71+ TEST_AUTH_HEADER ,
72+ "x-some-tenant-header" ,
73+ "v1" ,
74+ "grpc-trace-bin" ,
75+ "AAARf5ZpQwlN/8FVe1axOPlaAQIdRU/Y8j0LAgE" ),
76+ headersMap );
77+ }
6478
65- @ Override
66- public void fail (Status status ) {
67- Assertions .fail ("Failed in CallCredentials.MetadataApplier" );
68- }
69- });
79+ @ Override
80+ public void fail (Status status ) {
81+ Assertions .fail ("Failed in CallCredentials.MetadataApplier" );
82+ }
83+ });
7084 } finally {
7185 ctx .detach (previous );
7286 }
7387 }
74-
75-
7688}
0 commit comments