refactor out common code

also test with new probe_model

4 files changed, 167 insertions, 196 deletions

diff --git a/eval/src/tests/apps/analyze_onnx_model/analyze_onnx_model_test.cpp b/eval/src/tests/apps/analyze_onnx_model/analyze_onnx_model_test.cpp
index 2c1b2b21b9e..58cc7c22358 100644
--- a/eval/src/tests/apps/analyze_onnx_model/analyze_onnx_model_test.cpp
+++ b/eval/src/tests/apps/analyze_onnx_model/analyze_onnx_model_test.cpp
@@ -1,135 +1,37 @@
 // Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 
 #include <vespa/vespalib/testkit/test_kit.h>
-#include <vespa/vespalib/testkit/time_bomb.h>
-#include <vespa/vespalib/util/stringfmt.h>
 #include <vespa/vespalib/data/slime/slime.h>
-#include <vespa/vespalib/util/child_process.h>
-#include <vespa/vespalib/data/input.h>
-#include <vespa/vespalib/data/output.h>
-#include <vespa/vespalib/data/simple_buffer.h>
-#include <vespa/vespalib/util/size_literals.h>
 #include <vespa/eval/eval/test/test_io.h>
 
 using namespace vespalib;
 using namespace vespalib::eval::test;
-using vespalib::make_string_short::fmt;
-using vespalib::slime::JsonFormat;
-using vespalib::slime::Inspector;
 
 vespalib::string module_build_path("../../../../");
 vespalib::string binary = module_build_path + "src/apps/analyze_onnx_model/vespa-analyze-onnx-model";
 vespalib::string probe_cmd = binary + " --probe-types";
 
-std::string get_source_dir() {
+vespalib::string get_source_dir() {
     const char *dir = getenv("SOURCE_DIRECTORY");
     return (dir ? dir : ".");
 }
-std::string source_dir = get_source_dir();
-std::string guess_batch_model = source_dir + "/../../tensor/onnx_wrapper/guess_batch.onnx";
+vespalib::string source_dir = get_source_dir();
+vespalib::string probe_model = source_dir + "/../../tensor/onnx_wrapper/probe_model.onnx";
 
 //-----------------------------------------------------------------------------
 
-void read_until_eof(Input &input) {
-    for (auto mem = input.obtain(); mem.size > 0; mem = input.obtain()) {
-        input.evict(mem.size);
-    }
-}
-
-// Output adapter used to write to stdin of a child process
-class ChildIn : public Output {
-    ChildProcess &_child;
-    SimpleBuffer _output;
-public:
-    ChildIn(ChildProcess &child) : _child(child) {}
-    WritableMemory reserve(size_t bytes) override {
-        return _output.reserve(bytes);
-    }
-    Output &commit(size_t bytes) override {
-        _output.commit(bytes);
-        Memory buf = _output.obtain();
-        ASSERT_TRUE(_child.write(buf.data, buf.size));
-        _output.evict(buf.size);
-        return *this;
-    }
-};
-
-// Input adapter used to read from stdout of a child process
-class ChildOut : public Input {
-    ChildProcess &_child;
-    SimpleBuffer _input;
-public:
-    ChildOut(ChildProcess &child)
-      : _child(child)
-    {
-        EXPECT_TRUE(_child.running());
-        EXPECT_TRUE(!_child.failed());
-    }
-    Memory obtain() override {
-        if ((_input.get().size == 0) && !_child.eof()) {
-            WritableMemory buf = _input.reserve(4_Ki);
-            uint32_t res = _child.read(buf.data, buf.size);
-            ASSERT_TRUE((res > 0) || _child.eof());
-            _input.commit(res);
-        }
-        return _input.obtain();
-    }
-    Input &evict(size_t bytes) override {
-        _input.evict(bytes);
-        return *this;
-    }
-};
-
-//-----------------------------------------------------------------------------
-
-void dump_message(const char *prefix, const Slime &slime) {
-    SimpleBuffer buf;
-    slime::JsonFormat::encode(slime, buf, true);
-    auto str = buf.get().make_string();
-    fprintf(stderr, "%s%s\n", prefix, str.c_str());
-}
-
-class Server {
-private:
-    TimeBomb _bomb;
-    ChildProcess _child;
-    ChildIn _child_stdin;
-    ChildOut _child_stdout;
-public:
-    Server(vespalib::string cmd)
-      : _bomb(60),
-        _child(cmd.c_str()),
-        _child_stdin(_child),
-        _child_stdout(_child) {}
-    ~Server();
-    Slime invoke(const Slime &req) {
-        dump_message("request --> ", req);
-        write_compact(req, _child_stdin);
-        Slime reply;
-        ASSERT_TRUE(JsonFormat::decode(_child_stdout, reply));
-        dump_message("  reply <-- ", reply);
-        return reply;
-    }
-};
-Server::~Server() {
-    _child.close();
-    read_until_eof(_child_stdout);
-    ASSERT_TRUE(_child.wait());
-    ASSERT_TRUE(!_child.running());
-    ASSERT_TRUE(!_child.failed());
-}
-
-//-----------------------------------------------------------------------------
-
-TEST_F("require that output types can be probed", Server(probe_cmd)) {
+TEST_F("require that output types can be probed", ServerCmd(probe_cmd, true)) {
     Slime params;
     params.setObject();
-    params.get().setString("model", guess_batch_model);
+    params.get().setString("model", probe_model);
     params.get().setObject("inputs");
-    params["inputs"].setString("in1", "tensor<float>(x[3])");
-    params["inputs"].setString("in2", "tensor<float>(x[3])");
+    params["inputs"].setString("in1", "tensor<float>(x[2],y[3])");
+    params["inputs"].setString("in2", "tensor<float>(x[2],y[3])");
     Slime result = f1.invoke(params);
-    EXPECT_EQUAL(result["outputs"]["out"].asString().make_string(), vespalib::string("tensor<float>(d0[3])"));
+    EXPECT_EQUAL(result["outputs"].fields(), 3u);
+    EXPECT_EQUAL(result["outputs"]["out1"].asString().make_string(), vespalib::string("tensor<float>(d0[2],d1[3])"));
+    EXPECT_EQUAL(result["outputs"]["out2"].asString().make_string(), vespalib::string("tensor<float>(d0[2],d1[3])"));
+    EXPECT_EQUAL(result["outputs"]["out3"].asString().make_string(), vespalib::string("tensor<float>(d0[2],d1[3])"));
 }
 
 //-----------------------------------------------------------------------------
diff --git a/eval/src/tests/apps/eval_expr/eval_expr_test.cpp b/eval/src/tests/apps/eval_expr/eval_expr_test.cpp
index 077d3d2f35a..3834ea0e4cc 100644
--- a/eval/src/tests/apps/eval_expr/eval_expr_test.cpp
+++ b/eval/src/tests/apps/eval_expr/eval_expr_test.cpp
@@ -23,58 +23,6 @@ vespalib::string server_cmd = binary + " json-repl";
 
 //-----------------------------------------------------------------------------
 
-void read_until_eof(Input &input) {
-    for (auto mem = input.obtain(); mem.size > 0; mem = input.obtain()) {
-        input.evict(mem.size);
-    }
-}
-
-// Output adapter used to write to stdin of a child process
-class ChildIn : public Output {
-    ChildProcess &_child;
-    SimpleBuffer _output;
-public:
-    ChildIn(ChildProcess &child) : _child(child) {}
-    WritableMemory reserve(size_t bytes) override {
-        return _output.reserve(bytes);
-    }
-    Output &commit(size_t bytes) override {
-        _output.commit(bytes);
-        Memory buf = _output.obtain();
-        ASSERT_TRUE(_child.write(buf.data, buf.size));
-        _output.evict(buf.size);
-        return *this;
-    }
-};
-
-// Input adapter used to read from stdout of a child process
-class ChildOut : public Input {
-    ChildProcess &_child;
-    SimpleBuffer _input;
-public:
-    ChildOut(ChildProcess &child)
-      : _child(child)
-    {
-        EXPECT_TRUE(_child.running());
-        EXPECT_TRUE(!_child.failed());
-    }
-    Memory obtain() override {
-        if ((_input.get().size == 0) && !_child.eof()) {
-            WritableMemory buf = _input.reserve(4_Ki);
-            uint32_t res = _child.read(buf.data, buf.size);
-            ASSERT_TRUE((res > 0) || _child.eof());
-            _input.commit(res);
-        }
-        return _input.obtain();
-    }
-    Input &evict(size_t bytes) override {
-        _input.evict(bytes);
-        return *this;
-    }
-};
-
-//-----------------------------------------------------------------------------
-
 struct Result {
     vespalib::string error;
     vespalib::string result;
@@ -106,34 +54,9 @@ struct Result {
 };
 Result::~Result() = default;
 
-void dump_message(const char *prefix, const Slime &slime) {
-    SimpleBuffer buf;
-    slime::JsonFormat::encode(slime, buf, true);
-    auto str = buf.get().make_string();
-    fprintf(stderr, "%s%s\n", prefix, str.c_str());
-}
-
-class Server {
-private:
-    TimeBomb _bomb;
-    ChildProcess _child;
-    ChildIn _child_stdin;
-    ChildOut _child_stdout;
-public:
-    Server()
-      : _bomb(60),
-        _child(server_cmd.c_str()),
-        _child_stdin(_child),
-        _child_stdout(_child) {}
-    ~Server();
-    Slime invoke(const Slime &req) {
-        dump_message("request --> ", req);
-        write_compact(req, _child_stdin);
-        Slime reply;
-        ASSERT_TRUE(JsonFormat::decode(_child_stdout, reply));
-        dump_message("  reply <-- ", reply);
-        return reply;
-    }
+struct Server : public ServerCmd {
+    TimeBomb time_bomb;
+    Server() : ServerCmd(server_cmd, true), time_bomb(60) {}
     Result eval(const vespalib::string &expr, const vespalib::string &name = {}, bool verbose = false) {
         Slime req;
         auto &obj = req.setObject();
@@ -148,13 +71,6 @@ public:
         return {reply.get()};
     }
 };
-Server::~Server() {
-    _child.close();
-    read_until_eof(_child_stdout);
-    ASSERT_TRUE(_child.wait());
-    ASSERT_TRUE(!_child.running());
-    ASSERT_TRUE(!_child.failed());
-}
 
 //-----------------------------------------------------------------------------
 
diff --git a/eval/src/vespa/eval/eval/test/test_io.cpp b/eval/src/vespa/eval/eval/test/test_io.cpp
index fc7faef6555..ed7723f0748 100644
--- a/eval/src/vespa/eval/eval/test/test_io.cpp
+++ b/eval/src/vespa/eval/eval/test/test_io.cpp
@@ -5,6 +5,7 @@
 #include <vespa/vespalib/data/slime/slime.h>
 #include <vespa/vespalib/data/slime/json_format.h>
 #include <vespa/vespalib/util/size_literals.h>
+#include <filesystem>
 #include <unistd.h>
 #include <assert.h>
 
@@ -13,6 +14,8 @@ using vespalib::WritableMemory;
 using vespalib::slime::JsonFormat;
 using vespalib::slime::Cursor;
 
+namespace fs = std::filesystem;
+
 namespace vespalib::eval::test {
 
 //-----------------------------------------------------------------------------
@@ -63,6 +66,101 @@ StdOut::commit(size_t bytes)
 
 //-----------------------------------------------------------------------------
 
+ChildIn::ChildIn(ChildProcess &child)
+  : _child(child),
+    _output()
+{
+}
+
+WritableMemory
+ChildIn::reserve(size_t bytes)
+{
+    return _output.reserve(bytes);
+}
+
+Output &
+ChildIn::commit(size_t bytes)
+{
+    _output.commit(bytes);
+    Memory buf = _output.obtain();
+    REQUIRE(_child.write(buf.data, buf.size));
+    _output.evict(buf.size);
+    return *this;
+}
+
+//-----------------------------------------------------------------------------
+
+ChildOut::ChildOut(ChildProcess &child)
+  : _child(child),
+    _input()
+{
+    REQUIRE(_child.running());
+    REQUIRE(!_child.failed());
+}
+
+Memory
+ChildOut::obtain()
+{
+    if ((_input.get().size == 0) && !_child.eof()) {
+        WritableMemory buf = _input.reserve(4_Ki);
+        uint32_t res = _child.read(buf.data, buf.size);
+        REQUIRE((res > 0) || _child.eof());
+        _input.commit(res);
+    }
+    return _input.obtain();
+}
+
+Input &
+ChildOut::evict(size_t bytes)
+{
+    _input.evict(bytes);
+    return *this;
+}
+
+//-----------------------------------------------------------------------------
+
+void
+ServerCmd::maybe_dump_message(const char *prefix, const Slime &slime)
+{
+    if (_verbose) {
+        SimpleBuffer buf;
+        slime::JsonFormat::encode(slime, buf, false);
+        auto str = buf.get().make_string();
+        fprintf(stderr, "%s%s: %s", prefix, _basename.c_str(), str.c_str());
+    }
+}
+
+ServerCmd::ServerCmd(vespalib::string cmd, bool verbose)
+  : _child(cmd.c_str()),
+    _child_stdin(_child),
+    _child_stdout(_child),
+    _basename(fs::path(cmd).filename()),
+    _verbose(verbose)
+{
+}
+
+ServerCmd::~ServerCmd()
+{
+    _child.close();
+    read_until_eof(_child_stdout);
+    assert(_child.wait());
+    assert(!_child.running());
+    assert(!_child.failed());
+}
+
+Slime
+ServerCmd::invoke(const Slime &req)
+{
+    maybe_dump_message("request --> ", req);
+    write_compact(req, _child_stdin);
+    Slime reply;
+    REQUIRE(JsonFormat::decode(_child_stdout, reply));
+    maybe_dump_message("reply <-- ", reply);
+    return reply;
+}
+
+//-----------------------------------------------------------------------------
+
 bool
 LineReader::read_line(vespalib::string &line)
 {
@@ -96,6 +194,12 @@ bool look_for_eof(Input &input) {
     return true;
 }
 
+void read_until_eof(Input &input) {
+    for (auto mem = input.obtain(); mem.size > 0; mem = input.obtain()) {
+        input.evict(mem.size);
+    }
+}
+
 void write_compact(const Slime &slime, Output &out) {
     JsonFormat::encode(slime, out, true);
     out.reserve(1).data[0] = '\n';
diff --git a/eval/src/vespa/eval/eval/test/test_io.h b/eval/src/vespa/eval/eval/test/test_io.h
index b270b07bb30..cef6cd6ae40 100644
--- a/eval/src/vespa/eval/eval/test/test_io.h
+++ b/eval/src/vespa/eval/eval/test/test_io.h
@@ -9,6 +9,7 @@
 #include <vespa/vespalib/data/simple_buffer.h>
 #include <vespa/vespalib/data/slime/slime.h>
 #include <vespa/vespalib/util/size_literals.h>
+#include <vespa/vespalib/util/child_process.h>
 #include <functional>
 
 namespace vespalib::eval::test {
@@ -39,6 +40,49 @@ public:
 };
 
 /**
+ * Output adapter used to write to stdin of a child process.
+ **/
+class ChildIn : public Output {
+    ChildProcess &_child;
+    SimpleBuffer _output;
+public:
+    ChildIn(ChildProcess &child);
+    WritableMemory reserve(size_t bytes) override;
+    Output &commit(size_t bytes) override;
+};
+
+/**
+ * Input adapter used to read from stdout of a child process.
+ **/
+class ChildOut : public Input {
+    ChildProcess &_child;
+    SimpleBuffer _input;
+public:
+    ChildOut(ChildProcess &child);
+    Memory obtain() override;
+    Input &evict(size_t bytes) override;
+};
+
+/**
+ * A command run as a child process that acts as a server reading json
+ * from stdin and writing json to stdout.
+ **/
+class ServerCmd {
+private:
+    ChildProcess _child;
+    ChildIn _child_stdin;
+    ChildOut _child_stdout;
+    vespalib::string _basename;
+    bool _verbose;
+
+    void maybe_dump_message(const char *prefix, const Slime &slime);
+public:
+    ServerCmd(vespalib::string cmd, bool verbose);
+    ~ServerCmd();
+    Slime invoke(const Slime &req);
+};
+
+/**
  * Read one line at a time from an input
  **/
 class LineReader {
@@ -55,6 +99,11 @@ public:
 bool look_for_eof(Input &input);
 
 /**
+ * Read from the input until eof is reached (data is discarded).
+ **/
+void read_until_eof(Input &input);
+
+/**
  * Write a slime structure as compact json with a trailing newline.
  **/
 void write_compact(const Slime &slime, Output &out);