Merge pull request #22122 from vespa-engine/jonmv/fix-unit-test

Wait in mock destination session for RPC timeout

3 files changed, 35 insertions, 42 deletions

diff --git a/documentapi/src/test/java/com/yahoo/documentapi/messagebus/Destination.java b/documentapi/src/test/java/com/yahoo/documentapi/messagebus/Destination.java
index e08d9b14ee5..81eb237713f 100644
--- a/documentapi/src/test/java/com/yahoo/documentapi/messagebus/Destination.java
+++ b/documentapi/src/test/java/com/yahoo/documentapi/messagebus/Destination.java
@@ -4,9 +4,11 @@ package com.yahoo.documentapi.messagebus;
 import com.yahoo.document.DocumentRemove;
 import com.yahoo.documentapi.DocumentAccess;
 import com.yahoo.documentapi.DocumentAccessParams;
+import com.yahoo.documentapi.ProgressToken;
 import com.yahoo.documentapi.SyncParameters;
 import com.yahoo.documentapi.SyncSession;
 import com.yahoo.documentapi.local.LocalDocumentAccess;
+import com.yahoo.documentapi.messagebus.protocol.CreateVisitorReply;
 import com.yahoo.documentapi.messagebus.protocol.DocumentMessage;
 import com.yahoo.documentapi.messagebus.protocol.DocumentProtocol;
 import com.yahoo.documentapi.messagebus.protocol.GetDocumentMessage;
@@ -27,14 +29,18 @@ import com.yahoo.messagebus.network.Identity;
 import com.yahoo.messagebus.network.rpc.RPCNetworkParams;
 
 import java.util.Arrays;
+import java.util.List;
+import java.util.concurrent.Phaser;
 
 /**
  * Mock-up destination used for testing.
  *
- * @author <a href="mailto:einarmr@yahoo-inc.com">Einar M R Rosenvinge</a>
+ * @author Einar M R Rosenvinge
  */
 public class Destination implements MessageHandler {
 
+    final Phaser phaser = new Phaser(1);
+
     private final DestinationSession session;
     private final DocumentAccess access;
     private final SyncSession local;
@@ -46,7 +52,7 @@ public class Destination implements MessageHandler {
         params.setDocumentManagerConfigId(documentManagerConfigId);
         access = new LocalDocumentAccess(params);
         local = access.createSyncSession(new SyncParameters.Builder().build());
-        bus = new RPCMessageBus(Arrays.asList((Protocol)new DocumentProtocol(access.getDocumentTypeManager())),
+        bus = new RPCMessageBus(List.of(new DocumentProtocol(access.getDocumentTypeManager())),
                                 new RPCNetworkParams().setNumNetworkThreads(1)
                                         .setIdentity(new Identity("test/destination"))
                                         .setSlobrokConfigId(slobrokConfigId),
@@ -54,33 +60,35 @@ public class Destination implements MessageHandler {
         session = bus.getMessageBus().createDestinationSession("session", true, this);
     }
 
-    protected void sendReply(Reply reply) {
-        session.reply(reply);
-    }
-
     public void handleMessage(Message msg) {
+
+        phaser.arriveAndAwaitAdvance();
         Reply reply = ((DocumentMessage)msg).createReply();
         try {
             switch (msg.getType()) {
 
-            case DocumentProtocol.MESSAGE_GETDOCUMENT:
-                reply = new GetDocumentReply(local.get(((GetDocumentMessage)msg).getDocumentId()));
-                break;
+                case DocumentProtocol.MESSAGE_GETDOCUMENT:
+                    reply = new GetDocumentReply(local.get(((GetDocumentMessage)msg).getDocumentId()));
+                    break;
+
+                case DocumentProtocol.MESSAGE_PUTDOCUMENT:
+                    local.put(((PutDocumentMessage)msg).getDocumentPut());
+                    break;
 
-            case DocumentProtocol.MESSAGE_PUTDOCUMENT:
-                local.put(((PutDocumentMessage)msg).getDocumentPut());
-                break;
+                case DocumentProtocol.MESSAGE_REMOVEDOCUMENT:
+                    local.remove(new DocumentRemove(((RemoveDocumentMessage)msg).getDocumentId()));
+                    break;
 
-            case DocumentProtocol.MESSAGE_REMOVEDOCUMENT:
-                local.remove(new DocumentRemove(((RemoveDocumentMessage)msg).getDocumentId()));
-                break;
+                case DocumentProtocol.MESSAGE_UPDATEDOCUMENT:
+                    local.update(((UpdateDocumentMessage)msg).getDocumentUpdate());
+                    break;
 
-            case DocumentProtocol.MESSAGE_UPDATEDOCUMENT:
-                local.update(((UpdateDocumentMessage)msg).getDocumentUpdate());
-                break;
+                case DocumentProtocol.MESSAGE_CREATEVISITOR:
+                    ((CreateVisitorReply) reply).setLastBucket(ProgressToken.FINISHED_BUCKET);
+                    break;
 
-            default:
-                throw new UnsupportedOperationException("Unsupported message type '" + msg.getType() + "'.");
+                default:
+                    throw new UnsupportedOperationException("Unsupported message type '" + msg.getType() + "'.");
             }
         } catch (Exception e) {
             reply = new EmptyReply();
diff --git a/documentapi/src/test/java/com/yahoo/documentapi/messagebus/MessageBusDocumentApiTestCase.java b/documentapi/src/test/java/com/yahoo/documentapi/messagebus/MessageBusDocumentApiTestCase.java
index db7ab0ea238..44b600fe69f 100644
--- a/documentapi/src/test/java/com/yahoo/documentapi/messagebus/MessageBusDocumentApiTestCase.java
+++ b/documentapi/src/test/java/com/yahoo/documentapi/messagebus/MessageBusDocumentApiTestCase.java
@@ -12,6 +12,7 @@ import com.yahoo.documentapi.DocumentAccess;
 import com.yahoo.documentapi.DocumentOperationParameters;
 import com.yahoo.documentapi.ProgressToken;
 import com.yahoo.documentapi.Response;
+import com.yahoo.documentapi.Response.Outcome;
 import com.yahoo.documentapi.VisitorParameters;
 import com.yahoo.documentapi.VisitorSession;
 import com.yahoo.documentapi.messagebus.protocol.CreateVisitorReply;
@@ -70,7 +71,7 @@ public class MessageBusDocumentApiTestCase extends AbstractDocumentApiTestCase {
         params.setTraceLevel(9);
         access = new MessageBusDocumentAccess(params);
 
-        destination = new VisitableDestination(slobrokConfigId, params.getDocumentManagerConfigId());
+        destination = new Destination(slobrokConfigId, params.getDocumentManagerConfigId());
     }
 
     @After
@@ -80,25 +81,6 @@ public class MessageBusDocumentApiTestCase extends AbstractDocumentApiTestCase {
         slobrok.stop();
     }
 
-    private static class VisitableDestination extends Destination {
-        private VisitableDestination(String slobrokConfigId, String documentManagerConfigId) {
-            super(slobrokConfigId, documentManagerConfigId);
-        }
-
-        public void handleMessage(Message msg) {
-            if (msg.getType() == DocumentProtocol.MESSAGE_CREATEVISITOR) {
-                Reply reply = ((DocumentMessage)msg).createReply();
-                msg.swapState(reply);
-                CreateVisitorReply visitorReply = (CreateVisitorReply)reply;
-                visitorReply.setLastBucket(ProgressToken.FINISHED_BUCKET);
-                sendReply(reply);
-            } else {
-                super.handleMessage(msg);
-            }
-        }
-    }
-
-
     @Test
     public void requireThatVisitorSessionWorksWithMessageBus() throws ParseException, InterruptedException {
         VisitorParameters parameters = new VisitorParameters("id.user==1234");
@@ -119,17 +101,19 @@ public class MessageBusDocumentApiTestCase extends AbstractDocumentApiTestCase {
         AsyncSession session = access().createAsyncSession(new AsyncParameters());
         DocumentType type = access().getDocumentTypeManager().getDocumentType("music");
         Document doc1 = new Document(type, new DocumentId("id:ns:music::1"));
+
+        destination.phaser.register();
         assertTrue(session.put(new DocumentPut(doc1),
                                DocumentOperationParameters.parameters()
                                                           .withResponseHandler(result -> {
                                                               response.set(result);
                                                               latch.countDown();
                                                           })
-                                                          .withDeadline(Instant.now().minusSeconds(1)))
+                                                          .withDeadline(Instant.now().plusMillis(100)))
                           .isSuccess());
         assertTrue(latch.await(60, TimeUnit.SECONDS));
-        assertNotNull(response.get());
         assertEquals(Response.Outcome.TIMEOUT, response.get().outcome());
+        destination.phaser.arriveAndDeregister();
         session.destroy();
     }
 
diff --git a/documentapi/src/test/java/com/yahoo/documentapi/test/AbstractDocumentApiTestCase.java b/documentapi/src/test/java/com/yahoo/documentapi/test/AbstractDocumentApiTestCase.java
index a004ad368e0..001a5d284e2 100644
--- a/documentapi/src/test/java/com/yahoo/documentapi/test/AbstractDocumentApiTestCase.java
+++ b/documentapi/src/test/java/com/yahoo/documentapi/test/AbstractDocumentApiTestCase.java
@@ -173,4 +173,5 @@ public abstract class AbstractDocumentApiTestCase {
             latch.countDown();
         }
     }
+
 }