1 files changed, 376 insertions, 0 deletions
diff --git a/vespaclient-java/src/test/java/com/yahoo/vespaget/DocumentRetrieverTest.java b/vespaclient-java/src/test/java/com/yahoo/vespaget/DocumentRetrieverTest.java
new file mode 100644
index 00000000000..c3d3fcc71e9
--- /dev/null
+++ b/vespaclient-java/src/test/java/com/yahoo/vespaget/DocumentRetrieverTest.java
@@ -0,0 +1,376 @@
+// Copyright 2017 Yahoo Inc. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.vespaget;
+
+import com.fasterxml.jackson.core.JsonFactory;
+import com.fasterxml.jackson.core.JsonParseException;
+import com.fasterxml.jackson.core.JsonParser;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.yahoo.document.DataType;
+import com.yahoo.document.Document;
+import com.yahoo.document.DocumentId;
+import com.yahoo.documentapi.messagebus.MessageBusDocumentAccess;
+import com.yahoo.documentapi.messagebus.MessageBusParams;
+import com.yahoo.documentapi.messagebus.MessageBusSyncSession;
+import com.yahoo.documentapi.messagebus.loadtypes.LoadType;
+import com.yahoo.documentapi.messagebus.loadtypes.LoadTypeSet;
+import com.yahoo.documentapi.messagebus.protocol.DocumentProtocol;
+import com.yahoo.documentapi.messagebus.protocol.GetDocumentMessage;
+import com.yahoo.documentapi.messagebus.protocol.GetDocumentReply;
+import com.yahoo.messagebus.Error;
+import com.yahoo.messagebus.Reply;
+import com.yahoo.vespaclient.ClusterDef;
+import com.yahoo.vespaclient.ClusterList;
+
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.ExpectedException;
+import org.mockito.ArgumentMatcher;
+
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.io.PrintStream;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Map;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+import static org.mockito.Matchers.any;
+import static org.mockito.Matchers.argThat;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+/**
+ * Test class for {@link DocumentRetriever}
+ *
+ * @author bjorncs
+ */
+public class DocumentRetrieverTest {
+
+    public static final String DOC_ID_1 = "id:storage_test:document::1";
+    public static final String DOC_ID_2 = "id:storage_test:document::2";
+    public static final String DOC_ID_3 = "id:storage_test:document::3";
+
+    private final ByteArrayOutputStream outContent = new ByteArrayOutputStream();
+    private final ByteArrayOutputStream errContent = new ByteArrayOutputStream();
+
+    private DocumentAccessFactory mockedFactory;
+    private MessageBusDocumentAccess mockedDocumentAccess;
+    private MessageBusSyncSession mockedSession;
+    private PrintStream oldOut;
+    private PrintStream oldErr;
+
+    @Rule
+    public final ExpectedException exception = ExpectedException.none();
+
+    @Before
+    public void setUpStreams() {
+        oldOut = System.out;
+        oldErr = System.err;
+        System.setOut(new PrintStream(outContent));
+        System.setErr(new PrintStream(errContent));
+    }
+
+    @Before
+    public void prepareMessageBusMocks() {
+        this.mockedFactory = mock(DocumentAccessFactory.class);
+        this.mockedDocumentAccess = mock(MessageBusDocumentAccess.class);
+        this.mockedSession = mock(MessageBusSyncSession.class);
+        when(mockedFactory.createDocumentAccess(any())).thenReturn(mockedDocumentAccess);
+        when(mockedDocumentAccess.createSyncSession(any())).thenReturn(mockedSession);
+    }
+
+    @After
+    public void cleanUpStreams() {
+        System.setOut(oldOut);
+        System.setErr(oldErr);
+        outContent.reset();
+        errContent.reset();
+    }
+
+    private static ClientParameters.Builder createParameters() {
+        return new ClientParameters.Builder()
+                .setPriority(DocumentProtocol.Priority.NORMAL_2)
+                .setCluster("")
+                .setRoute("default")
+                .setConfigId("client")
+                .setFieldSet("[all]")
+                .setPrintIdsOnly(false)
+                .setHelp(false)
+                .setShowDocSize(false)
+                .setLoadTypeName("")
+                .setNoRetry(false)
+                .setTraceLevel(0)
+                .setTimeout(0)
+                .setDocumentIds(Collections.emptyIterator());
+    }
+
+    private static Iterator<String> asIterator(String... docIds) {
+        return Arrays.asList(docIds).iterator();
+    }
+
+    private static Reply createDocumentReply(String docId) {
+        return new GetDocumentReply(new Document(DataType.DOCUMENT, new DocumentId(docId)));
+    }
+
+    private void assertContainsDocument(String documentId) {
+        assertTrue(outContent.toString().contains(String.format(
+                "<document documenttype=\"document\" documentid=\"%s\"/>", documentId)));
+    }
+
+    private DocumentRetriever createDocumentRetriever(ClientParameters params) {
+        return createDocumentRetriever(params, new ClusterList());
+    }
+
+    private DocumentRetriever createDocumentRetriever(ClientParameters params, ClusterList clusterList) {
+        return new DocumentRetriever(
+                clusterList,
+                mockedFactory,
+                new LoadTypeSet(),
+                params);
+    }
+
+    @Test
+    public void testSendSingleMessage() throws DocumentRetrieverException {
+        ClientParameters params = createParameters()
+                .setDocumentIds(asIterator(DOC_ID_1))
+                .setPriority(DocumentProtocol.Priority.HIGH_1)
+                .setNoRetry(true)
+                .setLoadTypeName("loadtype")
+                .build();
+
+        when(mockedSession.syncSend(any())).thenReturn(createDocumentReply(DOC_ID_1));
+
+        LoadTypeSet loadTypeSet = new LoadTypeSet();
+        loadTypeSet.addLoadType(1, "loadtype", DocumentProtocol.Priority.HIGH_1);
+        DocumentRetriever documentRetriever = new DocumentRetriever(
+                new ClusterList(),
+                mockedFactory,
+                loadTypeSet,
+                params);
+        documentRetriever.retrieveDocuments();
+
+        verify(mockedSession, times(1)).syncSend(argThat(new ArgumentMatcher<GetDocumentMessage>() {
+            @Override
+            public boolean matches(Object o) {
+                GetDocumentMessage msg = (GetDocumentMessage) o;
+                return msg.getPriority().equals(DocumentProtocol.Priority.HIGH_1) &&
+                        !msg.getRetryEnabled() &&
+                        msg.getLoadType().equals(new LoadType(1, "loadtype", DocumentProtocol.Priority.HIGH_1));
+            }
+        }));
+        assertContainsDocument(DOC_ID_1);
+    }
+
+    @Test
+    public void testMultipleMessages() throws DocumentRetrieverException {
+        ClientParameters params = createParameters()
+                .setDocumentIds(asIterator(DOC_ID_1, DOC_ID_2, DOC_ID_3))
+                .build();
+
+        when(mockedSession.syncSend(any())).thenReturn(
+                createDocumentReply(DOC_ID_1),
+                createDocumentReply(DOC_ID_2),
+                createDocumentReply(DOC_ID_3));
+
+        DocumentRetriever documentRetriever = createDocumentRetriever(params);
+        documentRetriever.retrieveDocuments();
+
+        verify(mockedSession, times(3)).syncSend(any());
+        assertContainsDocument(DOC_ID_1);
+        assertContainsDocument(DOC_ID_2);
+        assertContainsDocument(DOC_ID_3);
+    }
+
+    @Test
+    public void testJsonOutput() throws DocumentRetrieverException, JsonParseException, IOException {
+        ClientParameters params = createParameters()
+                .setDocumentIds(asIterator(DOC_ID_1, DOC_ID_2, DOC_ID_3))
+                .setJsonOutput(true)
+                .build();
+
+        when(mockedSession.syncSend(any())).thenReturn(
+                createDocumentReply(DOC_ID_1),
+                createDocumentReply(DOC_ID_2),
+                createDocumentReply(DOC_ID_3));
+
+        DocumentRetriever documentRetriever = createDocumentRetriever(params);
+        documentRetriever.retrieveDocuments();
+
+        verify(mockedSession, times(3)).syncSend(any());
+        ObjectMapper m = new ObjectMapper();
+        @SuppressWarnings("unchecked")
+        List<Map<String, Object>> feed = m.readValue(outContent.toByteArray(), List.class);
+        assertEquals(DOC_ID_1, feed.get(0).get("id"));
+        assertEquals(DOC_ID_2, feed.get(1).get("id"));
+        assertEquals(DOC_ID_3, feed.get(2).get("id"));
+    }
+
+    @Test
+    public void testShutdownHook() throws DocumentRetrieverException {
+        ClientParameters params = createParameters()
+                .setDocumentIds(asIterator(DOC_ID_1))
+                .build();
+
+        when(mockedSession.syncSend(any())).thenReturn(createDocumentReply(DOC_ID_1));
+
+        DocumentRetriever documentRetriever = createDocumentRetriever(params);
+        documentRetriever.retrieveDocuments();
+        documentRetriever.shutdown();
+
+        verify(mockedSession, times(1)).destroy();
+        verify(mockedDocumentAccess, times(1)).shutdown();
+    }
+
+    @Test
+    public void testInvalidLoadType() throws DocumentRetrieverException {
+        exception.expect(DocumentRetrieverException.class);
+        exception.expectMessage("Loadtype with name 'undefinedloadtype' does not exist.\n");
+
+        ClientParameters params = createParameters()
+                .setLoadTypeName("undefinedloadtype")
+                .build();
+
+        DocumentRetriever documentRetriever = createDocumentRetriever(params);
+        documentRetriever.retrieveDocuments();
+    }
+
+    @Test
+    public void testClusterLookup() throws DocumentRetrieverException {
+        final String cluster = "storage", configId = "content/cluster.foo/storage",
+                expectedRoute = "[Storage:cluster=storage;clusterconfigid=content/cluster.foo/storage]";
+
+        ClientParameters params = createParameters()
+                .setCluster(cluster)
+                .build();
+
+        ClusterList clusterList = new ClusterList(Collections.singletonList(new ClusterDef(cluster, configId)));
+
+        DocumentRetriever documentRetriever = createDocumentRetriever(params, clusterList);
+        documentRetriever.retrieveDocuments();
+
+        verify(mockedFactory).createDocumentAccess(argThat(new ArgumentMatcher<MessageBusParams>() {
+            @Override
+            public boolean matches(Object o) {
+                return ((MessageBusParams) o).getRoute().equals(expectedRoute);
+            }
+        }));
+    }
+
+    @Test
+    public void testInvalidClusterName() throws DocumentRetrieverException {
+        exception.expect(DocumentRetrieverException.class);
+        exception.expectMessage("The Vespa cluster contains the content clusters storage, not invalidclustername. Please select a valid vespa cluster.");
+
+        ClientParameters params = createParameters()
+                .setCluster("invalidclustername")
+                .build();
+
+        ClusterList clusterList = new ClusterList(Collections.singletonList(new ClusterDef("storage", "content/cluster.foo/storage")));
+
+        DocumentRetriever documentRetriever = createDocumentRetriever(params, clusterList);
+        documentRetriever.retrieveDocuments();
+    }
+
+    @Test
+    public void testEmtpyClusterList() throws DocumentRetrieverException {
+        exception.expect(DocumentRetrieverException.class);
+        exception.expectMessage("The Vespa cluster does not have any content clusters declared.");
+
+        ClientParameters params = createParameters()
+                .setCluster("invalidclustername")
+                .build();
+
+        DocumentRetriever documentRetriever = createDocumentRetriever(params);
+        documentRetriever.retrieveDocuments();
+    }
+
+    @Test
+    public void testHandlingErrorFromMessageBus() throws DocumentRetrieverException {
+        ClientParameters params = createParameters()
+                .setDocumentIds(asIterator(DOC_ID_1))
+                .build();
+
+        Reply r = new GetDocumentReply(null);
+        r.addError(new Error(0, "Error message"));
+        when(mockedSession.syncSend(any())).thenReturn(r);
+
+        DocumentRetriever documentRetriever = createDocumentRetriever(params);
+        documentRetriever.retrieveDocuments();
+
+        assertTrue(errContent.toString().contains("Request failed"));
+    }
+
+    @Test
+    public void testShowDocSize() throws DocumentRetrieverException {
+        ClientParameters params = createParameters()
+                .setDocumentIds(asIterator(DOC_ID_1))
+                .setShowDocSize(true)
+                .build();
+
+        Document document = new Document(DataType.DOCUMENT, new DocumentId(DOC_ID_1));
+        when(mockedSession.syncSend(any())).thenReturn(new GetDocumentReply(document));
+
+        DocumentRetriever documentRetriever = createDocumentRetriever(params);
+        documentRetriever.retrieveDocuments();
+
+        assertTrue(outContent.toString().contains(String.format("Document size: %d bytes", document.getSerializedSize())));
+    }
+
+    @Test
+    public void testPrintIdOnly() throws DocumentRetrieverException {
+        ClientParameters params = createParameters()
+                .setDocumentIds(asIterator(DOC_ID_1))
+                .setPrintIdsOnly(true)
+                .build();
+
+        when(mockedSession.syncSend(any())).thenReturn(createDocumentReply(DOC_ID_1));
+
+        DocumentRetriever documentRetriever = createDocumentRetriever(params);
+        documentRetriever.retrieveDocuments();
+
+        assertEquals(DOC_ID_1 + "\n", outContent.toString());
+    }
+
+    @Test
+    public void testDocumentNotFound() throws DocumentRetrieverException {
+        ClientParameters params = createParameters()
+                .setDocumentIds(asIterator(DOC_ID_1))
+                .setPrintIdsOnly(true)
+                .build();
+
+        when(mockedSession.syncSend(any())).thenReturn(new GetDocumentReply(null));
+
+        DocumentRetriever documentRetriever = createDocumentRetriever(params);
+        documentRetriever.retrieveDocuments();
+
+        verify(mockedSession, times(1)).syncSend(any());
+        assertEquals(outContent.toString(), "Document not found.\n");
+    }
+
+    @Test
+    public void testTrace() throws DocumentRetrieverException {
+        final int traceLevel = 9;
+        ClientParameters params = createParameters()
+                .setDocumentIds(asIterator(DOC_ID_1))
+                .setTraceLevel(traceLevel)
+                .build();
+
+        GetDocumentReply reply = new GetDocumentReply(new Document(DataType.DOCUMENT, new DocumentId(DOC_ID_1)));
+        reply.getTrace().getRoot().addChild("childnode");
+        when(mockedSession.syncSend(any())).thenReturn(reply);
+
+        DocumentRetriever documentRetriever = createDocumentRetriever(params);
+        documentRetriever.retrieveDocuments();
+
+        verify(mockedSession, times(1)).setTraceLevel(traceLevel);
+        assertTrue(outContent.toString().contains("<trace>"));
+    }
+
+}