Added customization of GrpcTransport to JUnit4 & JUnit5 support

tests/common/src/main/java/tech/ydb/test/integration/YdbHelper.java

@@ -1,14 +1,24 @@
 package tech.ydb.test.integration;
 
 import tech.ydb.core.grpc.GrpcTransport;
+import tech.ydb.core.grpc.GrpcTransportBuilder;
 
 /**
  *
  * @author Aleksandr Gorshenin
  */
 public interface YdbHelper extends AutoCloseable {
+    @FunctionalInterface
+    interface TransportCustomizer {
+        GrpcTransportBuilder apply(GrpcTransportBuilder builder);
+    }
+
     GrpcTransport createTransport();
 
+    default GrpcTransport createTransport(TransportCustomizer customizer) {
+        return createTransport();
+    }
+
     String endpoint();
     String database();
 

tests/common/src/main/java/tech/ydb/test/integration/docker/DockerHelperFactory.java

@@ -35,9 +35,17 @@ public YdbHelper createHelper() {
         return new YdbHelper() {
             @Override
             public GrpcTransport createTransport() {
+                return createTransport(null);
+            }
+
+            @Override
+            public GrpcTransport createTransport(TransportCustomizer customiser) {
                 GrpcTransportBuilder builder = GrpcTransport.forEndpoint(endpoint(), container.database());
                 if (env.ydbUseTls()) {
-                    builder.withSecureConnection(container.pemCert());
+                    builder = builder.withSecureConnection(container.pemCert());
+                }
+                if (customiser != null) {
+                    builder = customiser.apply(builder);
                 }
                 return builder.build();
             }

tests/common/src/main/java/tech/ydb/test/integration/docker/ProxedDockerHelperFactory.java

@@ -8,6 +8,7 @@
 import org.testcontainers.utility.TestcontainersConfiguration;
 
 import tech.ydb.core.grpc.GrpcTransport;
+import tech.ydb.core.grpc.GrpcTransportBuilder;
 import tech.ydb.test.integration.YdbEnvironment;
 import tech.ydb.test.integration.YdbHelper;
 import tech.ydb.test.integration.YdbHelperFactory;
@@ -42,7 +43,16 @@ public YdbHelper createHelper() {
         return new YdbHelper() {
             @Override
             public GrpcTransport createTransport() {
-                return GrpcTransport.forEndpoint(endpoint(), container.database()).build();
+                return createTransport(null);
+            }
+
+            @Override
+            public GrpcTransport createTransport(TransportCustomizer customizer) {
+                GrpcTransportBuilder builder = GrpcTransport.forEndpoint(endpoint(), container.database());
+                if (customizer != null) {
+                    builder = customizer.apply(builder);
+                }
+                return builder.build();
             }
 
             @Override

tests/common/src/main/java/tech/ydb/test/integration/external/ExternalHelperFactory.java

@@ -29,13 +29,22 @@ public YdbHelper createHelper() {
         return new YdbHelper() {
             @Override
             public GrpcTransport createTransport() {
+                return createTransport(null);
+            }
+
+            @Override
+            public GrpcTransport createTransport(TransportCustomizer customizer) {
                 GrpcTransportBuilder builder = GrpcTransport.forEndpoint(endpoint(), database());
 
                 if (authToken() != null) {
-                    builder.withAuthProvider(new TokenAuthProvider(authToken()));
+                    builder = builder.withAuthProvider(new TokenAuthProvider(authToken()));
                 }
                 if (useTls()) {
-                    builder.withSecureConnection(pemCert());
+                    builder = builder.withSecureConnection(pemCert());
+                }
+
+                if (customizer != null) {
+                    builder = customizer.apply(builder);
                 }
 
                 return builder.build();

tests/common/src/main/java/tech/ydb/test/integration/utils/ProxyYdbHelper.java

@@ -24,6 +24,11 @@ public GrpcTransport createTransport() {
         return checked().createTransport();
     }
 
+    @Override
+    public GrpcTransport createTransport(TransportCustomizer customizer) {
+        return checked().createTransport(customizer);
+    }
+
     @Override
     public String endpoint() {
         return checked().endpoint();

tests/common/src/test/java/tech/ydb/test/integration/DockerHelperFactoryTest.java

@@ -147,7 +147,7 @@ public void defaultDockerContainerTests() {
             Assert.assertNull("check helper auth token", helper.authToken());
             Assert.assertArrayEquals("check helper database", container.pemCert, helper.pemCert());
 
-            try (GrpcTransport transport = helper.createTransport()) {
+            try (GrpcTransport transport = helper.createTransport(builder -> builder.withApplicationName("test-app"))) {
                 Assert.assertEquals("/local", transport.getDatabase());
                 Assert.assertTrue(transport.unaryCall(null, null, null).join().isSuccess());
             }

tests/common/src/test/java/tech/ydb/test/integration/GrpcTransportMock.java

@@ -24,6 +24,11 @@ public void setup(String database) {
         staticGrpcTransport.when(() -> GrpcTransport.forEndpoint(Mockito.any(), Mockito.any()))
                 .thenReturn(builderMock);
 
+        Mockito.when(builderMock.withAuthProvider(Mockito.any())).thenReturn(builderMock);
+        Mockito.when(builderMock.withApplicationName(Mockito.any())).thenReturn(builderMock);
+        Mockito.when(builderMock.withSecureConnection(Mockito.any())).thenReturn(builderMock);
+        Mockito.when(builderMock.withSecureConnection()).thenReturn(builderMock);
+
         Mockito.when(builderMock.build()).thenReturn(transportMock);
 
         Mockito.when(transportMock.getDatabase()).thenReturn(database);

tests/common/src/test/java/tech/ydb/test/integration/ProxyDockerTest.java

@@ -52,6 +52,28 @@ public void testProxedDocker() throws InvalidProtocolBufferException {
                 Assert.assertNull(helper.pemCert());
                 Assert.assertFalse(helper.useTls());
             }
+
+            try (GrpcTransport transport = helper.createTransport(builder -> builder.withApplicationName("proxed"))) {
+                GrpcRequestSettings settings = GrpcRequestSettings.newBuilder().build();
+                SchemeOperationProtos.DescribePathRequest request = SchemeOperationProtos.DescribePathRequest
+                        .newBuilder()
+                        .setPath(helper.database())
+                        .build();
+
+                SchemeOperationProtos.DescribePathResponse response = transport.unaryCall(
+                        SchemeServiceGrpc.getDescribePathMethod(), settings, request
+                ).join().getValue();
+
+                Assert.assertTrue(response.getOperation().getReady());
+
+                SchemeOperationProtos.DescribePathResult result = response.getOperation().getResult()
+                        .unpack(SchemeOperationProtos.DescribePathResult.class);
+
+                Assert.assertEquals(helper.database(), "/" + result.getSelf().getName());
+                Assert.assertNull(helper.authToken());
+                Assert.assertNull(helper.pemCert());
+                Assert.assertFalse(helper.useTls());
+            }
         }
     }
 }

tests/common/src/test/java/tech/ydb/test/integration/YdbHelperFactoryTest.java

@@ -138,6 +138,11 @@ public void externalAuthInstanceTest() {
                 Assert.assertEquals("/token", transport.getDatabase());
                 Assert.assertTrue(transport.unaryCall(null, null, null).join().isSuccess());
             }
+
+            try (GrpcTransport transport = helper.createTransport(builder -> builder.withApplicationName("test"))) {
+                Assert.assertEquals("/token", transport.getDatabase());
+                Assert.assertTrue(transport.unaryCall(null, null, null).join().isSuccess());
+            }
         }
     }
 

tests/junit4-support/src/main/java/tech/ydb/test/junit4/GrpcTransportRule.java

@@ -22,19 +22,26 @@ public class GrpcTransportRule extends ProxyGrpcTransport implements TestRule {
     private static final Logger logger = LoggerFactory.getLogger(GrpcTransportRule.class);
 
     private final AtomicReference<GrpcTransport> proxy;
+    private final YdbHelper.TransportCustomizer customizer;
     private final boolean skipOnUnavailable;
 
-    private GrpcTransportRule(AtomicReference<GrpcTransport> proxy, boolean skipOnUnavailable) {
+    private GrpcTransportRule(AtomicReference<GrpcTransport> proxy, YdbHelper.TransportCustomizer customizer,
+            boolean skipOnUnavailable) {
         this.proxy = proxy;
+        this.customizer = customizer;
         this.skipOnUnavailable = skipOnUnavailable;
     }
 
     public GrpcTransportRule() {
-        this(new AtomicReference<>(), true);
+        this(new AtomicReference<>(), null, true);
+    }
+
+    public GrpcTransportRule withGrpcTransportCustomizer(YdbHelper.TransportCustomizer customizer) {
+        return new GrpcTransportRule(proxy, customizer, skipOnUnavailable);
     }
 
     public GrpcTransportRule failIfUnavailable() {
-        return new GrpcTransportRule(proxy, false);
+        return new GrpcTransportRule(proxy, customizer, false);
     }
 
     @Override
@@ -60,7 +67,7 @@ public void evaluate() throws Throwable {
 
                 logger.debug("create ydb helper for test {}", path);
                 try (YdbHelper helper = factory.createHelper()) {
-                    try (GrpcTransport transport = helper.createTransport()) {
+                    try (GrpcTransport transport = helper.createTransport(customizer)) {
                         proxy.set(transport);
                         base.evaluate();
                         proxy.set(null);

tests/junit5-support/src/main/java/tech/ydb/test/junit5/GrpcTransportExtension.java

@@ -34,7 +34,11 @@ private GrpcTransportExtension(Holder holder, boolean skipOnUnavailable) {
     }
 
     public GrpcTransportExtension() {
-        this(new Holder(), true);
+        this(new Holder(null), true);
+    }
+
+    public GrpcTransportExtension withGrpcTransportCustomizer(YdbHelper.TransportCustomizer customizer) {
+        return new GrpcTransportExtension(new Holder(customizer), skipOnUnavailable);
     }
 
     public GrpcTransportExtension failIfUnavailable() {
@@ -82,11 +86,16 @@ public void afterEach(ExtensionContext ec) throws Exception {
 
     private static class Holder {
         private final Lock holderLock = new ReentrantLock();
+        private final YdbHelper.TransportCustomizer customizer;
 
         private YdbHelper helper = null;
         private GrpcTransport transport = null;
         private ExtensionContext context = null;
 
+        Holder(YdbHelper.TransportCustomizer customizer) {
+            this.customizer = customizer;
+        }
+
         public void before(ExtensionContext ec) {
             holderLock.lock();
             try {
@@ -107,7 +116,7 @@ public void before(ExtensionContext ec) {
                     }
 
                     logger.debug("create ydb helper for path {}", path);
-                    transport = helper.createTransport();
+                    transport = helper.createTransport(customizer);
                 }
             } finally {
                 holderLock.unlock();