diff options
Diffstat (limited to 'container-search/src/test/java/com/yahoo/search/dispatch/rpc/MockClient.java')
-rw-r--r-- | container-search/src/test/java/com/yahoo/search/dispatch/rpc/MockClient.java | 144 |
1 files changed, 144 insertions, 0 deletions
diff --git a/container-search/src/test/java/com/yahoo/search/dispatch/rpc/MockClient.java b/container-search/src/test/java/com/yahoo/search/dispatch/rpc/MockClient.java new file mode 100644 index 00000000000..f9b628e594a --- /dev/null +++ b/container-search/src/test/java/com/yahoo/search/dispatch/rpc/MockClient.java @@ -0,0 +1,144 @@ +// Copyright 2017 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root. +package com.yahoo.search.dispatch.rpc; + +import com.yahoo.compress.CompressionType; +import com.yahoo.compress.Compressor; +import com.yahoo.document.GlobalId; +import com.yahoo.document.idstring.IdIdString; +import com.yahoo.prelude.fastsearch.FastHit; +import com.yahoo.search.Result; +import com.yahoo.search.dispatch.rpc.Client; +import com.yahoo.search.dispatch.rpc.RpcFillInvoker; +import com.yahoo.search.dispatch.rpc.RpcSearchInvoker; +import com.yahoo.slime.ArrayTraverser; +import com.yahoo.slime.BinaryFormat; +import com.yahoo.slime.Cursor; +import com.yahoo.slime.Inspector; +import com.yahoo.slime.Slime; + +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +/** + * @author bratseth + */ +public class MockClient implements Client { + + private final Map<DocsumKey, Map<String, Object>> docsums = new HashMap<>(); + private final Compressor compressor = new Compressor(); + private boolean malfunctioning = false; + private Result searchResult; + + /** Set to true to cause this to produce an error instead of a regular response */ + public void setMalfunctioning(boolean malfunctioning) { this.malfunctioning = malfunctioning; } + + @Override + public NodeConnection createConnection(String hostname, int port) { + return new MockNodeConnection(hostname, port); + } + + @Override + public void getDocsums(List<FastHit> hitsContext, NodeConnection node, CompressionType compression, + int uncompressedSize, byte[] compressedSlime, RpcFillInvoker.GetDocsumsResponseReceiver responseReceiver, + double timeoutSeconds) { + if (malfunctioning) { + responseReceiver.receive(GetDocsumsResponseOrError.fromError("Malfunctioning")); + return; + } + + Inspector request = BinaryFormat.decode(compressor.decompress(compressedSlime, compression, uncompressedSize)).get(); + String docsumClass = request.field("class").asString(); + List<Map<String, Object>> docsumsToReturn = new ArrayList<>(); + request.field("gids").traverse((ArrayTraverser)(index, gid) -> { + GlobalId docId = new GlobalId(gid.asData()); + docsumsToReturn.add(docsums.get(new DocsumKey(node.toString(), docId, docsumClass))); + }); + Slime responseSlime = new Slime(); + Cursor root = responseSlime.setObject(); + Cursor docsums = root.setArray("docsums"); + for (Map<String, Object> docsumFields : docsumsToReturn) { + Cursor docsumItem = docsums.addObject(); + Cursor docsum = docsumItem.setObject("docsum"); + for (Map.Entry<String, Object> field : docsumFields.entrySet()) { + if (field.getValue() instanceof Integer) + docsum.setLong(field.getKey(), (Integer)field.getValue()); + else if (field.getValue() instanceof String) + docsum.setString(field.getKey(), (String)field.getValue()); + else + throw new RuntimeException(); + } + } + byte[] slimeBytes = BinaryFormat.encode(responseSlime); + Compressor.Compression compressionResult = compressor.compress(compression, slimeBytes); + GetDocsumsResponse response = new GetDocsumsResponse(compressionResult.type().getCode(), slimeBytes.length, + compressionResult.data(), hitsContext); + responseReceiver.receive(GetDocsumsResponseOrError.fromResponse(response)); + } + + @Override + public void search(NodeConnection node, CompressionType compression, int uncompressedLength, byte[] compressedPayload, + RpcSearchInvoker responseReceiver, double timeoutSeconds) { + if (malfunctioning) { + responseReceiver.receive(SearchResponseOrError.fromError("Malfunctioning")); + return; + } + + if(searchResult == null) { + responseReceiver.receive(SearchResponseOrError.fromError("No result defined")); + return; + } + var payload = ProtobufSerialization.serializeResult(searchResult); + var compressionResult = compressor.compress(compression, payload); + var response = new SearchResponse(compressionResult.type().getCode(), payload.length, compressionResult.data()); + responseReceiver.receive(SearchResponseOrError.fromResponse(response)); + } + + public void setDocsumReponse(String nodeId, int docId, String docsumClass, Map<String, Object> docsumValues) { + docsums.put(new DocsumKey(nodeId, globalIdFrom(docId), docsumClass), docsumValues); + } + + public GlobalId globalIdFrom(int hitId) { + return new GlobalId(new IdIdString("", "test", "", String.valueOf(hitId))); + } + + private static class MockNodeConnection implements Client.NodeConnection { + + private final String hostname; + + public MockNodeConnection(String hostname, int port) { + this.hostname = hostname; + } + + @Override + public void close() { } + + @Override + public String toString() { return hostname; } + + } + + private static class DocsumKey { + + private final String internalKey; + + public DocsumKey(String nodeId, GlobalId docId, String docsumClass) { + internalKey = docsumClass + "." + nodeId + "." + docId; + } + + @Override + public int hashCode() { return internalKey.hashCode(); } + + @Override + public boolean equals(Object other) { + if ( ! (other instanceof DocsumKey)) return false; + return ((DocsumKey)other).internalKey.equals(this.internalKey); + } + + @Override + public String toString() { return internalKey; } + + } + +} |