Allow running expectations against stubbed functions

Co-authored-by: hydraz <urn@semi.works>

src/test/resources/test-rom/mcfly.lua

@@ -27,18 +27,52 @@ local function check(func, arg, ty, val)
     end
 end
 
+--- A stub - wraps a value within a a table,
+local stub_mt = {}
+stub_mt.__index = stub_mt
+
+--- Revert this stub, restoring the previous value.
+--
+-- Note, a stub can only be reverted once.
+function stub_mt:revert()
+    if not self.active then return end
+
+    self.active = false
+    rawset(self.stubbed_in, self.key, self.original)
+end
+
 local active_stubs = {}
 
---- Stub a global variable with a specific value
+--- Stub a table entry with a new value.
 --
--- @tparam string var The variable to stub
--- @param value The value to stub it with
-local function stub(tbl, var, value)
+-- @tparam table
+-- @tparam string key The variable to stub
+-- @param value The value to stub it with. If this is a function, one can use
+-- the various stub expectation methods to determine what it was called with.
+-- @treturn Stub The resulting stub
+local function stub(tbl, key, value)
     check('stub', 1, 'table', tbl)
-    check('stub', 2, 'string', var)
+    check('stub', 2, 'string', key)
 
-    table.insert(active_stubs, { tbl = tbl, var = var, value = tbl[var] })
-    rawset(tbl, var, value)
+    local stub = setmetatable({
+        active = true,
+        stubbed_in = tbl,
+        key = key,
+        original = rawget(tbl, key),
+    }, stub_mt)
+
+    if type(value) == "function" then
+        local arguments, delegate = {}, value
+        stub.arguments = arguments
+        value = function(...)
+            arguments[#arguments + 1] = table.pack(...)
+            return delegate(value)
+        end
+    end
+
+    table.insert(active_stubs, stub)
+    rawset(tbl, key, value)
+    return stub
 end
 
 --- Capture the current global state of the computer
@@ -57,10 +91,7 @@ end
 
 --- Restore the global state of the computer to a previous version
 local function pop_state(state)
-    for i = #active_stubs, 1, -1 do
-        local stub = active_stubs[i]
-        rawset(stub.tbl, stub.var, stub.value)
-    end
+    for i = #active_stubs, 1, -1 do active_stubs[i]:revert() end
 
     active_stubs = state.stubs
 
@@ -210,6 +241,16 @@ local function matches(eq, exact, left, right)
     return true
 end
 
+local function pairwise_equal(left, right)
+    if left.n ~= right.n then return false end
+
+    for i = 1, left.n do
+        if left[i] ~= right[i] then return false end
+    end
+
+    return true
+end
+
 --- Assert that this expectation is structurally equivalent to
 -- the provided object.
 --
@@ -236,6 +277,59 @@ function expect_mt:matches(value)
     return self
 end
 
+--- Assert that this stub was called a specific number of times.
+--
+-- @tparam[opt] number The exact number of times the function must be called.
+-- If not given just require the function to be called at least once.
+-- @raises If this function was not called the expected number of times.
+function expect_mt:called(times)
+    if getmetatable(self.value) ~= stub_mt or self.value.arguments == nil then
+        fail(("Expected stubbed function, got %s"):format(type(self.value)))
+    end
+
+    local called = #self.value.arguments
+
+    if times == nil then
+        if called == 0 then
+            fail("Expected stub to be called\nbut it was not.")
+        end
+    else
+        check('stub', 1, 'number', times)
+        if called ~= times then
+            fail(("Expected stub to be called %d times\nbut was called %d times."):format(times, called))
+        end
+    end
+
+    return self
+end
+
+--- Assert that this stub was called with a set of arguments
+--
+-- Arguments are compared using exact equality.
+function expect_mt:called_with(...)
+    if getmetatable(self.value) ~= stub_mt or self.value.arguments == nil then
+        fail(("Expected stubbed function, got %s"):format(type(self.value)))
+    end
+
+    local exp_args = table.pack(...)
+    local actual_args = self.value.arguments
+    for i = 1, #actual_args do
+        if pairwise_equal(actual_args[i], exp_args) then return self end
+    end
+
+    local head = ("Expected stub to be called with %s\nbut was"):format(format(exp_args))
+    if #actual_args == 0 then
+        fail(head .. " not called at all")
+    elseif #actual_args == 1 then
+        fail(("%s called with %s."):format(head, format(actual_args[1])))
+    else
+        local lines = { head .. " called with:" }
+        for i = 1, #actual_args do lines[i + 1] = " - " .. format(actual_args[i]) end
+
+        fail(table.concat(lines, "\n"))
+    end
+end
+
 local expect = setmetatable( {
     --- Construct an expectation on the error message calling this function
     -- produces