From 18ef620679245e3161618ba4bed760489647b1c0 Mon Sep 17 00:00:00 2001
From: Davin McCall <davmac@davmac.org>
Date: Wed, 1 Jan 2020 10:53:13 +1000
Subject: [PATCH] Improve unit test infrastructure - allow control of write()

In the unit test infrastructure, allow behaviour of the write() call
(i.e. bp_sys::write()) to be overridden with a custom handler.
---
 src/tests/test-bpsys.cc                | 33 ++++++++++++++++++++------
 src/tests/test-includes/baseproc-sys.h | 22 +++++++++++++++++
 src/tests/test-includes/dinit.h        |  7 +++++-
 3 files changed, 54 insertions(+), 8 deletions(-)

diff --git a/src/tests/test-bpsys.cc b/src/tests/test-bpsys.cc
index 21e9b46..fc0a08d 100644
--- a/src/tests/test-bpsys.cc
+++ b/src/tests/test-bpsys.cc
@@ -1,6 +1,7 @@
 #include <vector>
 #include <utility>
 #include <algorithm>
+#include <memory>
 #include <map>
 
 #include <cstdlib>
@@ -35,8 +36,8 @@ class read_cond : public std::vector<read_result>
 // map of fd to read results to supply for reads of that fd
 std::map<int,read_cond> read_data;
 
-// map of data written to each fd
-std::map<int,std::vector<char>> written_data;
+// map of fd to the handler for writes to that fd
+std::map<int, std::unique_ptr<bp_sys::write_handler>> write_hndlr_map;
 
 } // anon namespace
 
@@ -47,18 +48,33 @@ pid_t last_forked_pid = 1;  // last forked process id (incremented each 'fork')
 
 // Test helper methods:
 
+void init_bpsys()
+{
+    write_hndlr_map[0] = std::unique_ptr<bp_sys::write_handler> { new default_write_handler() };
+    write_hndlr_map[1] = std::unique_ptr<bp_sys::write_handler> { new default_write_handler() };
+    write_hndlr_map[2] = std::unique_ptr<bp_sys::write_handler> { new default_write_handler() };
+}
+
 // Allocate a file descriptor
 int allocfd()
+{
+    return allocfd(new default_write_handler());
+}
+
+int allocfd(write_handler *whndlr)
 {
     auto f = std::find(usedfds.begin(), usedfds.end(), false);
     if (f == usedfds.end()) {
         int r = usedfds.size();
         usedfds.push_back(true);
+        write_hndlr_map[r] = std::unique_ptr<bp_sys::write_handler>(whndlr);
         return r;
     }
 
     *f = true;
-    return f - usedfds.begin();
+    auto r = f - usedfds.begin();
+    write_hndlr_map[r] = std::unique_ptr<bp_sys::write_handler>(whndlr);
+    return r;
 }
 
 // Supply data to be returned by read()
@@ -80,7 +96,10 @@ void set_blocking(int fd)
 // retrieve data written via write()
 void extract_written_data(int fd, std::vector<char> &data)
 {
-	data = std::move(written_data[fd]);
+    auto &whandler = write_hndlr_map[fd];
+    if (whandler == nullptr) abort();
+    default_write_handler *dwhndlr = static_cast<default_write_handler *>(whandler.get());
+	data = std::move(dwhndlr->data);
 }
 
 
@@ -96,8 +115,10 @@ int pipe2(int fds[2], int flags)
 int close(int fd)
 {
     if (size_t(fd) >= usedfds.size()) abort();
+    if (! usedfds[fd]) abort();
 
     usedfds[fd] = false;
+    write_hndlr_map.erase(fd);
     return 0;
 }
 
@@ -145,9 +166,7 @@ ssize_t read(int fd, void *buf, size_t count)
 
 ssize_t write(int fd, const void *buf, size_t count)
 {
-	std::vector<char> &wd = written_data[fd];
-	wd.insert(wd.end(), (char *)buf, (char *)buf + count);
-	return count;
+    return write_hndlr_map[fd]->write(fd, buf, count);
 }
 
 ssize_t writev(int fd, const struct iovec *iov, int iovcnt)
diff --git a/src/tests/test-includes/baseproc-sys.h b/src/tests/test-includes/baseproc-sys.h
index c8e2441..2d29807 100644
--- a/src/tests/test-includes/baseproc-sys.h
+++ b/src/tests/test-includes/baseproc-sys.h
@@ -12,10 +12,32 @@
 
 namespace bp_sys {
 
+class write_handler
+{
+public:
+    virtual ssize_t write(int fd, const void *buf, size_t count) = 0;
+    virtual ~write_handler() { }
+};
+
+class default_write_handler : public write_handler
+{
+public:
+    std::vector<char> data;
+
+    virtual ssize_t write(int fd, const void *buf, size_t count)
+    {
+        data.insert(data.end(), (char *)buf, (char *)buf + count);
+        return count;
+    }
+};
+
 // Test helper functions:
 
+void init_bpsys();
+
 // allocate a file descriptor
 int allocfd();
+int allocfd(write_handler *hndlr);
 
 void supply_read_data(int fd, std::vector<char> &data);
 void supply_read_data(int fd, std::vector<char> &&data);
diff --git a/src/tests/test-includes/dinit.h b/src/tests/test-includes/dinit.h
index bf6eb68..22804fc 100644
--- a/src/tests/test-includes/dinit.h
+++ b/src/tests/test-includes/dinit.h
@@ -47,7 +47,11 @@ class eventloop_t
     {
         auto i = regd_fd_watchers.find(fd);
         if (i != regd_fd_watchers.end()) {
-            i->second->fd_event(*this, fd, events);
+            fd_watcher *watcher = i->second;
+            dasynq::rearm r = watcher->fd_event(*this, fd, events);
+            if (r == dasynq::rearm::REMOVE) {
+                watcher->deregister(*this);
+            }
         }
     }
 
@@ -115,6 +119,7 @@ class eventloop_t
         {
             loop.regd_fd_watchers.erase(watched_fd);
             watched_fd = -1;
+            watch_removed();
         }
 
         virtual rearm fd_event(eventloop_t & loop, int fd, int flags) = 0;
-- 
2.25.1