1 files changed, 47 insertions, 14 deletions

diff --git a/tests/functional/repl_characterization/test-session.cc b/tests/functional/repl_characterization/test-session.cc
index 50e27e58c..52179a372 100644
--- a/tests/functional/repl_characterization/test-session.cc
+++ b/tests/functional/repl_characterization/test-session.cc
@@ -1,4 +1,5 @@
 #include <iostream>
+#include <span>
 #include <unistd.h>
 
 #include "test-session.hh"
@@ -21,14 +22,17 @@ RunningProcess RunningProcess::start(std::string executable, Strings args)
 
     // This is separate from runProgram2 because we have different IO requirements
     pid_t pid = startProcess([&]() {
-        if (dup2(procStdout.writeSide.get(), STDOUT_FILENO) == -1)
+        if (dup2(procStdout.writeSide.get(), STDOUT_FILENO) == -1) {
             throw SysError("dupping stdout");
-        if (dup2(procStdin.readSide.get(), STDIN_FILENO) == -1)
+        }
+        if (dup2(procStdin.readSide.get(), STDIN_FILENO) == -1) {
             throw SysError("dupping stdin");
+        }
         procStdin.writeSide.close();
         procStdout.readSide.close();
-        if (dup2(STDOUT_FILENO, STDERR_FILENO) == -1)
+        if (dup2(STDOUT_FILENO, STDERR_FILENO) == -1) {
             throw SysError("dupping stderr");
+        }
         execv(executable.c_str(), stringsToCharPtrs(args).data());
         throw SysError("exec did not happen");
     });
@@ -44,7 +48,8 @@ RunningProcess RunningProcess::start(std::string executable, Strings args)
 }
 
 [[gnu::unused]]
-std::ostream & operator<<(std::ostream & os, ReplOutputParser::State s)
+std::ostream &
+operator<<(std::ostream & os, ReplOutputParser::State s)
 {
     switch (s) {
     case ReplOutputParser::State::Prompt:
@@ -91,8 +96,7 @@ bool ReplOutputParser::feed(char c)
     return false;
 }
 
-/** Waits for the prompt and then returns if a prompt was found */
-bool TestSession::waitForPrompt()
+bool TestSession::readOutThen(ReadOutThenCallback cb)
 {
     std::vector<char> buf(1024);
 
@@ -106,38 +110,67 @@ bool TestSession::waitForPrompt()
             return false;
         }
 
+        switch (cb(std::span(buf.data(), res))) {
+        case ReadOutThenCallbackResult::Stop:
+            return true;
+        case ReadOutThenCallbackResult::Continue:
+            continue;
+        }
+    }
+}
+
+bool TestSession::waitForPrompt()
+{
+    bool notEof = readOutThen([&](std::span<char> s) -> ReadOutThenCallbackResult {
         bool foundPrompt = false;
-        for (ssize_t i = 0; i < res; ++i) {
+
+        for (auto ch : s) {
             // foundPrompt = foundPrompt || outputParser.feed(buf[i]);
             bool wasEaten = true;
-            eater.feed(buf[i], [&](char c) {
+            eater.feed(ch, [&](char c) {
                 wasEaten = false;
-                foundPrompt = outputParser.feed(buf[i]) || foundPrompt;
+                foundPrompt = outputParser.feed(ch) || foundPrompt;
 
                 outLog.push_back(c);
             });
 
             if constexpr (DEBUG_REPL_PARSER) {
-                std::cerr << "raw " << MaybeHexEscapedChar{buf[i]} << (wasEaten ? " [eaten]" : "") << "\n";
+                std::cerr << "raw " << MaybeHexEscapedChar{ch} << (wasEaten ? " [eaten]" : "") << "\n";
             }
         }
 
-        if (foundPrompt) {
-            return true;
+        return foundPrompt ? ReadOutThenCallbackResult::Stop : ReadOutThenCallbackResult::Continue;
+    });
+
+    return notEof;
+}
+
+void TestSession::wait()
+{
+    readOutThen([&](std::span<char> s) {
+        for (auto ch : s) {
+            eater.feed(ch, [&](char c) {
+                outputParser.feed(c);
+                outLog.push_back(c);
+            });
         }
-    }
+        // just keep reading till we hit eof
+        return ReadOutThenCallbackResult::Continue;
+    });
 }
 
 void TestSession::close()
 {
     proc.procStdin.close();
+    wait();
     proc.procStdout.close();
 }
 
 void TestSession::runCommand(std::string command)
 {
-    if constexpr (DEBUG_REPL_PARSER)
+    if constexpr (DEBUG_REPL_PARSER) {
         std::cerr << "runCommand " << command << "\n";
+    }
     command += "\n";
     // We have to feed a newline into the output parser, since Nix might not
     // give us a newline before a prompt in all cases (it might clear line