From 7b2f01d45ba636d3755ba26cfa5a4b75f93e9844 Mon Sep 17 00:00:00 2001
From: lethosor <lethosor@gmail.com>
Date: Fri, 9 Apr 2021 00:26:33 -0400
Subject: [PATCH] Add initial mock.patch() implementation for tests

---
 ci/test.lua                     |  2 ++
 library/lua/test_util/mock.lua  | 52 +++++++++++++++++++++++++++++++++
 test/library/test_util/mock.lua | 45 ++++++++++++++++++++++++++++
 3 files changed, 99 insertions(+)
 create mode 100644 library/lua/test_util/mock.lua
 create mode 100644 test/library/test_util/mock.lua

diff --git a/ci/test.lua b/ci/test.lua
index 5c12463f2..706560351 100644
--- a/ci/test.lua
+++ b/ci/test.lua
@@ -3,6 +3,7 @@
 
 local expect = require 'test_util.expect'
 local json = require 'json'
+local mock = require 'test_util.mock'
 local script = require 'gui.script'
 local utils = require 'utils'
 
@@ -187,6 +188,7 @@ local function build_test_env()
             mode = 'none',
         },
         expect = {},
+        mock = mock,
         delay = delay,
         require = clean_require,
         reqscript = clean_reqscript,
diff --git a/library/lua/test_util/mock.lua b/library/lua/test_util/mock.lua
new file mode 100644
index 000000000..1e4609f10
--- /dev/null
+++ b/library/lua/test_util/mock.lua
@@ -0,0 +1,52 @@
+local mock = mkmodule('test_util.mock')
+
+
+--[[
+Usage:
+    patch(table, key, value, callback)
+    patch({
+        [{table, key}] = value,
+        [{table2, key2}] = value2
+    }, callback)
+]]
+function mock.patch(...)
+    local args = {...}
+    if #args == 4 then
+        args = {{
+            [{args[1], args[2]}] = args[3]
+        }, args[4]}
+    end
+    if #args ~= 2 then
+        error('expected 2 or 4 arguments')
+    end
+
+    local callback = args[2]
+    local patches = {}
+    for k, v in pairs(args[1]) do
+        local p = {
+            table = k[1],
+            key = k[2],
+            new_value = v,
+        }
+        p.old_value = p.table[p.key]
+        -- no-op to ensure that the value can be restored by the finalizer below
+        p.table[p.key] = p.table[p.key]
+        table.insert(patches, p)
+    end
+
+    dfhack.with_finalize(
+        function()
+            for _, p in ipairs(patches) do
+                p.table[p.key] = p.old_value
+            end
+        end,
+        function()
+            for _, p in ipairs(patches) do
+                p.table[p.key] = p.new_value
+            end
+            callback()
+        end
+    )
+end
+
+return mock
diff --git a/test/library/test_util/mock.lua b/test/library/test_util/mock.lua
new file mode 100644
index 000000000..f43fcbd65
--- /dev/null
+++ b/test/library/test_util/mock.lua
@@ -0,0 +1,45 @@
+local mock = require('test_util.mock')
+
+local test_table = {
+    func1 = function()
+        return 1
+    end,
+    func2 = function()
+        return 2
+    end,
+}
+
+function test.patch_single()
+    expect.eq(test_table.func1(), 1)
+    mock.patch(test_table, 'func1', function() return 3 end, function()
+        expect.eq(test_table.func1(), 3)
+    end)
+    expect.eq(test_table.func1(), 1)
+end
+
+function test.patch_single_nested()
+    expect.eq(test_table.func1(), 1)
+    mock.patch(test_table, 'func1', function() return 3 end, function()
+        expect.eq(test_table.func1(), 3)
+        mock.patch(test_table, 'func1', function() return 5 end, function()
+            expect.eq(test_table.func1(), 5)
+        end)
+        expect.eq(test_table.func1(), 3)
+    end)
+    expect.eq(test_table.func1(), 1)
+end
+
+function test.patch_multiple()
+    expect.eq(test_table.func1(), 1)
+    expect.eq(test_table.func2(), 2)
+    mock.patch({
+        [{test_table, 'func1'}] = function() return 3 end,
+        [{test_table, 'func2'}] = function() return 4 end,
+    }, function()
+        expect.eq(test_table.func1(), 3)
+        expect.eq(test_table.func2(), 4)
+    end)
+    expect.eq(test_table.func1(), 1)
+    expect.eq(test_table.func2(), 2)
+end
+