Fixed test flapping

Author: alex268

query/src/test/java/tech/ydb/query/impl/GrpcTestInterceptor.java

@@ -21,14 +21,14 @@
  * @author Aleksandr Gorshenin
  */
 public class GrpcTestInterceptor implements Consumer<ManagedChannelBuilder<?>>, ClientInterceptor {
-    private final Queue<Status> nextStatus = new ConcurrentLinkedQueue<>();
+    private volatile Queue<Status> overrideQueue = new ConcurrentLinkedQueue<>();
 
     public void reset() {
-        nextStatus.clear();
+        overrideQueue = new ConcurrentLinkedQueue<>();
     }
 
-    public void addNextStatus(Status status) {
-        nextStatus.add(status);
+    public void addOverrideStatus(Status status) {
+        overrideQueue.add(status);
     }
 
     @Override
@@ -39,15 +39,17 @@ public void accept(ManagedChannelBuilder<?> t) {
     @Override
     public <ReqT, RespT> ClientCall<ReqT, RespT> interceptCall(
             MethodDescriptor<ReqT, RespT> method, CallOptions callOptions, Channel next) {
-        return new ProxyClientCall<>(next, method, callOptions);
+        return new ProxyClientCall<>(next, overrideQueue.poll(), method, callOptions);
     }
 
-    private class ProxyClientCall<ReqT, RespT> extends ClientCall<ReqT, RespT> {
+    private static class ProxyClientCall<ReqT, RespT> extends ClientCall<ReqT, RespT> {
         private final ClientCall<ReqT, RespT> delegate;
+        private final Status overrided;
 
-        private ProxyClientCall(Channel channel, MethodDescriptor<ReqT, RespT> method,
+        private ProxyClientCall(Channel channel, Status overrided, MethodDescriptor<ReqT, RespT> method,
                 CallOptions callOptions) {
             this.delegate = channel.newCall(method, callOptions);
+            this.overrided = overrided;
         }
 
         @Override
@@ -110,8 +112,7 @@ public void onMessage(RespT message) {
 
             @Override
             public void onClose(Status status, Metadata trailers) {
-                Status next = nextStatus.poll();
-                delegate.onClose(next != null ? next : status, trailers);
+                delegate.onClose(overrided != null ? overrided : status, trailers);
             }
 
             @Override

query/src/test/java/tech/ydb/query/impl/QueryClientTest.java

@@ -94,7 +94,7 @@ public void sessionExecuteQueryTest() {
         QuerySession s1 = getSession();
         String id1 = s1.getId();
 
-        grpcInterceptor.addNextStatus(io.grpc.Status.UNAVAILABLE);
+        grpcInterceptor.addOverrideStatus(io.grpc.Status.UNAVAILABLE);
 
         Result<QueryInfo> res = s1.createQuery("SELECT 1 + 2", TxMode.NONE).execute().join();
         Assert.assertEquals(StatusCode.TRANSPORT_UNAVAILABLE, res.getStatus().getCode());

table/src/test/java/tech/ydb/table/integration/GrpcTestInterceptor.java

@@ -21,14 +21,14 @@
  * @author Aleksandr Gorshenin
  */
 public class GrpcTestInterceptor implements Consumer<ManagedChannelBuilder<?>>, ClientInterceptor {
-    private final Queue<Status> nextStatus = new ConcurrentLinkedQueue<>();
+    private volatile Queue<Status> overrideQueue = new ConcurrentLinkedQueue<>();
 
     public void reset() {
-        nextStatus.clear();
+        overrideQueue = new ConcurrentLinkedQueue<>();
     }
 
-    public void addNextStatus(Status status) {
-        nextStatus.add(status);
+    public void addOverrideStatus(Status status) {
+        overrideQueue.add(status);
     }
 
     @Override
@@ -39,15 +39,17 @@ public void accept(ManagedChannelBuilder<?> t) {
     @Override
     public <ReqT, RespT> ClientCall<ReqT, RespT> interceptCall(
             MethodDescriptor<ReqT, RespT> method, CallOptions callOptions, Channel next) {
-        return new ProxyClientCall<>(next, method, callOptions);
+        return new ProxyClientCall<>(next, overrideQueue.poll(), method, callOptions);
     }
 
-    private class ProxyClientCall<ReqT, RespT> extends ClientCall<ReqT, RespT> {
+    private static class ProxyClientCall<ReqT, RespT> extends ClientCall<ReqT, RespT> {
         private final ClientCall<ReqT, RespT> delegate;
+        private final Status overrided;
 
-        private ProxyClientCall(Channel channel, MethodDescriptor<ReqT, RespT> method,
+        private ProxyClientCall(Channel channel, Status overrided, MethodDescriptor<ReqT, RespT> method,
                 CallOptions callOptions) {
             this.delegate = channel.newCall(method, callOptions);
+            this.overrided = overrided;
         }
 
         @Override
@@ -110,8 +112,7 @@ public void onMessage(RespT message) {
 
             @Override
             public void onClose(Status status, Metadata trailers) {
-                Status next = nextStatus.poll();
-                delegate.onClose(next != null ? next : status, trailers);
+                delegate.onClose(overrided != null ? overrided : status, trailers);
             }
 
             @Override

table/src/test/java/tech/ydb/table/integration/TableClientTest.java

@@ -97,7 +97,7 @@ public void sessionExecuteDataQueryTest() {
         Session s1 = getSession();
         String id1 = s1.getId();
 
-        grpcInterceptor.addNextStatus(io.grpc.Status.UNAVAILABLE);
+        grpcInterceptor.addOverrideStatus(io.grpc.Status.UNAVAILABLE);
 
         Result<DataQueryResult> res = s1.executeDataQuery("SELECT 1 + 2", TxControl.snapshotRo()).join();
         Assert.assertEquals(StatusCode.TRANSPORT_UNAVAILABLE, res.getStatus().getCode());
@@ -128,7 +128,7 @@ public void sessionExecuteScanQueryTest() {
         Session s1 = getSession();
         String id1 = s1.getId();
 
-        grpcInterceptor.addNextStatus(io.grpc.Status.UNAVAILABLE);
+        grpcInterceptor.addOverrideStatus(io.grpc.Status.UNAVAILABLE);
 
         Status res = s1.executeScanQuery("SELECT 1 + 2", Params.empty(), settings).start(rsr -> {}).join();
         Assert.assertEquals(StatusCode.TRANSPORT_UNAVAILABLE, res.getCode());