1 files changed, 84 insertions, 33 deletions

diff --git a/tests/unit/libutil/serialise.cc b/tests/unit/libutil/serialise.cc
index 95ae43115..78882ad2c 100644
--- a/tests/unit/libutil/serialise.cc
+++ b/tests/unit/libutil/serialise.cc
@@ -2,30 +2,47 @@
 #include "error.hh"
 #include "fmt.hh"
 #include "pos-table.hh"
+#include "generator.hh"
 #include "ref.hh"
 #include "types.hh"
 
+#include <concepts>
+#include <cstdint>
+#include <initializer_list>
 #include <limits.h>
 #include <gtest/gtest.h>
 
 #include <numeric>
+#include <stdexcept>
+#include <string_view>
+#include <type_traits>
 
 namespace nix {
 
-TEST(Sink, uint64_t)
+// don't deduce the type of `val` for added insurance.
+template<typename T>
+static std::string toWire(const std::type_identity_t<T> & val)
 {
-    StringSink s;
-    s << 42;
-    ASSERT_EQ(s.s, std::string({42, 0, 0, 0, 0, 0, 0, 0}));
+    std::string result;
+    auto g = [] (const auto & val) -> WireFormatGenerator { co_yield val; }(val);
+    while (auto buffer = g.next()) {
+        result.append(buffer->data(), buffer->size());
+    }
+    return result;
 }
 
-TEST(Sink, string_view)
+TEST(WireFormatGenerator, uint64_t)
 {
-    StringSink s;
-    s << "";
+    auto s = toWire<uint64_t>(42);
+    ASSERT_EQ(s, std::string({42, 0, 0, 0, 0, 0, 0, 0}));
+}
+
+TEST(WireFormatGenerator, string_view)
+{
+    auto s = toWire<std::string_view>("");
     // clang-format off
     ASSERT_EQ(
-        s.s,
+        s,
         std::string({
             // length
             0, 0, 0, 0, 0, 0, 0, 0,
@@ -34,11 +51,10 @@ TEST(Sink, string_view)
     );
     // clang-format on
 
-    s = {};
-    s << "test";
+    s = toWire<std::string_view>("test");
     // clang-format off
     ASSERT_EQ(
-        s.s,
+        s,
         std::string({
             // length
             4, 0, 0, 0, 0, 0, 0, 0,
@@ -50,11 +66,10 @@ TEST(Sink, string_view)
     );
     // clang-format on
 
-    s = {};
-    s << "longer string";
+    s = toWire<std::string_view>("longer string");
     // clang-format off
     ASSERT_EQ(
-        s.s,
+        s,
         std::string({
             // length
             13, 0, 0, 0, 0, 0, 0, 0,
@@ -67,13 +82,12 @@ TEST(Sink, string_view)
     // clang-format on
 }
 
-TEST(Sink, StringSet)
+TEST(WireFormatGenerator, StringSet)
 {
-    StringSink s;
-    s << StringSet{};
+    auto s = toWire<StringSet>({});
     // clang-format off
     ASSERT_EQ(
-        s.s,
+        s,
         std::string({
             // length
             0, 0, 0, 0, 0, 0, 0, 0,
@@ -82,11 +96,10 @@ TEST(Sink, StringSet)
     );
     // clang-format on
 
-    s = {};
-    s << StringSet{"a", ""};
+    s = toWire<StringSet>({"a", ""});
     // clang-format off
     ASSERT_EQ(
-        s.s,
+        s,
         std::string({
             // length
             2, 0, 0, 0, 0, 0, 0, 0,
@@ -99,13 +112,12 @@ TEST(Sink, StringSet)
     // clang-format on
 }
 
-TEST(Sink, Strings)
+TEST(WireFormatGenerator, Strings)
 {
-    StringSink s;
-    s << Strings{};
+    auto s = toWire<Strings>({});
     // clang-format off
     ASSERT_EQ(
-        s.s,
+        s,
         std::string({
             // length
             0, 0, 0, 0, 0, 0, 0, 0,
@@ -114,11 +126,10 @@ TEST(Sink, Strings)
     );
     // clang-format on
 
-    s = {};
-    s << Strings{"a", ""};
+    s = toWire<Strings>({"a", ""});
     // clang-format off
     ASSERT_EQ(
-        s.s,
+        s,
         std::string({
             // length
             2, 0, 0, 0, 0, 0, 0, 0,
@@ -131,23 +142,22 @@ TEST(Sink, Strings)
     // clang-format on
 }
 
-TEST(Sink, Error)
+TEST(WireFormatGenerator, Error)
 {
     PosTable pt;
     auto o = pt.addOrigin(Pos::String{make_ref<std::string>("test")}, 4);
 
-    StringSink s;
-    s << Error{ErrorInfo{
+    auto s = toWire<Error>(Error{ErrorInfo{
         .level = lvlInfo,
         .msg = HintFmt("foo"),
         .pos = pt[pt.add(o, 1)],
         .traces = {{.pos = pt[pt.add(o, 2)], .hint = HintFmt("b %1%", "foo")}},
-    }};
+    }});
     // NOTE position of the error and all traces are ignored
     // by the wire format
     // clang-format off
     ASSERT_EQ(
-        s.s,
+        s,
         std::string({
             5, 0, 0, 0, 0, 0, 0, 0, 'E', 'r', 'r', 'o', 'r', 0, 0, 0,
             3, 0, 0, 0, 0, 0, 0, 0,
@@ -163,4 +173,45 @@ TEST(Sink, Error)
     // clang-format on
 }
 
+TEST(WireFormatGenerator, exampleMessage)
+{
+    auto gen = []() -> WireFormatGenerator {
+        std::set<std::string> foo{"a", "longer string", ""};
+        co_yield 42;
+        co_yield foo;
+        co_yield std::string_view("test");
+        co_yield true;
+    }();
+
+    std::vector<char> full;
+    while (auto s = gen.next()) {
+        full.insert(full.end(), s->begin(), s->end());
+    }
+
+    ASSERT_EQ(
+        full,
+        (std::vector<char>{
+            // clang-format off
+            // 42
+            42, 0, 0, 0, 0, 0, 0, 0,
+            // foo
+            3, 0, 0, 0, 0, 0, 0, 0,
+            /// ""
+            0, 0, 0, 0, 0, 0, 0, 0,
+            /// a
+            1, 0, 0, 0, 0, 0, 0, 0,
+            'a', 0, 0, 0, 0, 0, 0, 0,
+            /// longer string
+            13, 0, 0, 0, 0, 0, 0, 0,
+            'l', 'o', 'n', 'g', 'e', 'r', ' ', 's', 't', 'r', 'i', 'n', 'g', 0, 0, 0,
+            // foo done
+            // test
+            4, 0, 0, 0, 0, 0, 0, 0,
+            't', 'e', 's', 't', 0, 0, 0, 0,
+            // true
+            1, 0, 0, 0, 0, 0, 0, 0,
+            //clang-format on
+            }));
+}
+
 }