From 23e9f20ebc1b0d013bd3f95a9463a9fe92a2cb44 Mon Sep 17 00:00:00 2001
From: Tor Brede Vekterli <vekterli@yahooinc.com>
Date: Mon, 5 Dec 2022 13:56:34 +0100
Subject: Add tensor short form output option to vespa-visit tool

Specified with `--tensor-short-form`. No single-char option alias,
as short form output will be the default on Vespa 9 and we're running
out of usable option characters for this tool anyway.
---
 .../yahoo/vespavisit/StdOutVisitorHandlerTest.java | 56 ++++++++++++++++++++--
 1 file changed, 53 insertions(+), 3 deletions(-)

(limited to 'vespaclient-java/src/test/java/com/yahoo/vespavisit/StdOutVisitorHandlerTest.java')

diff --git a/vespaclient-java/src/test/java/com/yahoo/vespavisit/StdOutVisitorHandlerTest.java b/vespaclient-java/src/test/java/com/yahoo/vespavisit/StdOutVisitorHandlerTest.java
index f109b0ad56b..a2e9f91d503 100644
--- a/vespaclient-java/src/test/java/com/yahoo/vespavisit/StdOutVisitorHandlerTest.java
+++ b/vespaclient-java/src/test/java/com/yahoo/vespavisit/StdOutVisitorHandlerTest.java
@@ -1,7 +1,18 @@
 // Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package com.yahoo.vespavisit;
 
+import com.yahoo.document.Document;
+import com.yahoo.document.DocumentPut;
+import com.yahoo.document.DocumentType;
+import com.yahoo.document.TensorDataType;
+import com.yahoo.document.datatypes.TensorFieldValue;
+import com.yahoo.documentapi.AckToken;
+import com.yahoo.documentapi.VisitorControlSession;
 import com.yahoo.documentapi.VisitorDataHandler;
+import com.yahoo.documentapi.messagebus.protocol.PutDocumentMessage;
+import com.yahoo.tensor.Tensor;
+import com.yahoo.tensor.TensorType;
+import org.junit.jupiter.api.Test;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.MethodSource;
 
@@ -9,6 +20,7 @@ import java.io.ByteArrayOutputStream;
 import java.io.PrintStream;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.mockito.Mockito.mock;
 
 /**
  * @author bjorncs
@@ -30,7 +42,7 @@ public class StdOutVisitorHandlerTest {
         initStdOutVisitorHandlerTest(jsonOutput);
         ByteArrayOutputStream out = new ByteArrayOutputStream();
         StdOutVisitorHandler visitorHandler =
-                new StdOutVisitorHandler(/*printIds*/true, false, false, false, false, false, 0, jsonOutput, new PrintStream(out, true));
+                new StdOutVisitorHandler(/*printIds*/true, false, false, false, false, false, 0, jsonOutput, false, new PrintStream(out, true));
         VisitorDataHandler dataHandler = visitorHandler.getDataHandler();
         dataHandler.onDone();
         String output = out.toString();
@@ -43,11 +55,49 @@ public class StdOutVisitorHandlerTest {
         initStdOutVisitorHandlerTest(jsonOutput);
         ByteArrayOutputStream out = new ByteArrayOutputStream();
         StdOutVisitorHandler visitorHandler =
-                new StdOutVisitorHandler(/*printIds*/false, false, false, false, false, false, 0, jsonOutput, new PrintStream(out, true));
+                new StdOutVisitorHandler(/*printIds*/false, false, false, false, false, false, 0, jsonOutput, false, new PrintStream(out, true));
         VisitorDataHandler dataHandler = visitorHandler.getDataHandler();
         dataHandler.onDone();
         String expectedOutput = jsonOutput ? "[]" : "";
         String output = out.toString().trim();
         assertEquals(expectedOutput, output);
     }
-}
\ No newline at end of file
+
+    void do_test_json_tensor_fields_can_be_output_in_short_or_long_form(boolean tensorShortForm, String expectedOutput) {
+        var docType = new DocumentType("foo");
+        docType.addField("bar", TensorDataType.getTensor(TensorType.fromSpec("tensor(x[3])")));
+        var doc = new Document(docType, "id:baz:foo::tensor-stuff");
+        doc.setFieldValue("bar", new TensorFieldValue(Tensor.from("tensor(x[3]):[1,2,3]")));
+        var putMsg = new PutDocumentMessage(new DocumentPut(doc));
+
+        var out = new ByteArrayOutputStream();
+        var visitorHandler = new StdOutVisitorHandler(/*printIds*/false, false, false, false, false, false,
+                                                      0, true, tensorShortForm, new PrintStream(out, true));
+        var dataHandler = visitorHandler.getDataHandler();
+        var controlSession = mock(VisitorControlSession.class);
+        var ackToken = mock(AckToken.class);
+        dataHandler.setSession(controlSession);
+        dataHandler.onMessage(putMsg, ackToken);
+        dataHandler.onDone();
+
+        String output = out.toString().trim();
+        assertEquals(expectedOutput, output);
+    }
+
+    @Test
+    void json_tensor_fields_can_be_output_in_long_form() {
+        var expectedOutput = """
+        [
+        {"id":"id:baz:foo::tensor-stuff","fields":{"bar":{"cells":[{"address":{"x":"0"},"value":1.0},{"address":{"x":"1"},"value":2.0},{"address":{"x":"2"},"value":3.0}]}}}]""";
+        do_test_json_tensor_fields_can_be_output_in_short_or_long_form(false, expectedOutput);
+    }
+
+    @Test
+    void json_tensor_fields_can_be_output_in_short_form() {
+        var expectedOutput = """
+        [
+        {"id":"id:baz:foo::tensor-stuff","fields":{"bar":{"type":"tensor(x[3])","values":[1.0,2.0,3.0]}}}]""";
+        do_test_json_tensor_fields_can_be_output_in_short_or_long_form(true, expectedOutput);
+    }
+
+}
-- 
cgit v1.2.3