diff --git a/tests/common/src/main/java/tech/ydb/test/integration/YdbHelper.java b/tests/common/src/main/java/tech/ydb/test/integration/YdbHelper.java index 55e3f506a..65ce311aa 100644 --- a/tests/common/src/main/java/tech/ydb/test/integration/YdbHelper.java +++ b/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(); diff --git a/tests/common/src/main/java/tech/ydb/test/integration/docker/DockerHelperFactory.java b/tests/common/src/main/java/tech/ydb/test/integration/docker/DockerHelperFactory.java index 6688abf25..8994b270d 100644 --- a/tests/common/src/main/java/tech/ydb/test/integration/docker/DockerHelperFactory.java +++ b/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(); } diff --git a/tests/common/src/main/java/tech/ydb/test/integration/docker/ProxedDockerHelperFactory.java b/tests/common/src/main/java/tech/ydb/test/integration/docker/ProxedDockerHelperFactory.java index 099a90d9b..f2d96f3f5 100644 --- a/tests/common/src/main/java/tech/ydb/test/integration/docker/ProxedDockerHelperFactory.java +++ b/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 diff --git a/tests/common/src/main/java/tech/ydb/test/integration/external/ExternalHelperFactory.java b/tests/common/src/main/java/tech/ydb/test/integration/external/ExternalHelperFactory.java index f1e64614d..821560807 100644 --- a/tests/common/src/main/java/tech/ydb/test/integration/external/ExternalHelperFactory.java +++ b/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(); diff --git a/tests/common/src/main/java/tech/ydb/test/integration/utils/ProxyYdbHelper.java b/tests/common/src/main/java/tech/ydb/test/integration/utils/ProxyYdbHelper.java index 23872dfbf..0b5fc50e5 100644 --- a/tests/common/src/main/java/tech/ydb/test/integration/utils/ProxyYdbHelper.java +++ b/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(); diff --git a/tests/common/src/test/java/tech/ydb/test/integration/DockerHelperFactoryTest.java b/tests/common/src/test/java/tech/ydb/test/integration/DockerHelperFactoryTest.java index 97662b24e..3ca61bc29 100644 --- a/tests/common/src/test/java/tech/ydb/test/integration/DockerHelperFactoryTest.java +++ b/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()); } diff --git a/tests/common/src/test/java/tech/ydb/test/integration/GrpcTransportMock.java b/tests/common/src/test/java/tech/ydb/test/integration/GrpcTransportMock.java index e4e4f0b78..7d8faa6fd 100644 --- a/tests/common/src/test/java/tech/ydb/test/integration/GrpcTransportMock.java +++ b/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); diff --git a/tests/common/src/test/java/tech/ydb/test/integration/ProxyDockerTest.java b/tests/common/src/test/java/tech/ydb/test/integration/ProxyDockerTest.java index 7d4ea7fdd..ab90be1b9 100644 --- a/tests/common/src/test/java/tech/ydb/test/integration/ProxyDockerTest.java +++ b/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()); + } } } } diff --git a/tests/common/src/test/java/tech/ydb/test/integration/YdbHelperFactoryTest.java b/tests/common/src/test/java/tech/ydb/test/integration/YdbHelperFactoryTest.java index 02631e6d8..4bfb242e2 100644 --- a/tests/common/src/test/java/tech/ydb/test/integration/YdbHelperFactoryTest.java +++ b/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()); + } } } diff --git a/tests/junit4-support/src/main/java/tech/ydb/test/junit4/GrpcTransportRule.java b/tests/junit4-support/src/main/java/tech/ydb/test/junit4/GrpcTransportRule.java index cc3fbdd55..a7852d644 100644 --- a/tests/junit4-support/src/main/java/tech/ydb/test/junit4/GrpcTransportRule.java +++ b/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 proxy; + private final YdbHelper.TransportCustomizer customizer; private final boolean skipOnUnavailable; - private GrpcTransportRule(AtomicReference proxy, boolean skipOnUnavailable) { + private GrpcTransportRule(AtomicReference 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); diff --git a/tests/junit5-support/src/main/java/tech/ydb/test/junit5/GrpcTransportExtension.java b/tests/junit5-support/src/main/java/tech/ydb/test/junit5/GrpcTransportExtension.java index 909043449..0a1e988bf 100644 --- a/tests/junit5-support/src/main/java/tech/ydb/test/junit5/GrpcTransportExtension.java +++ b/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();