Add balancer.lua unit tests

rootfs/etc/nginx/lua/balancer.lua

@@ -48,6 +48,9 @@ end
 
 local function balance()
   local backend = get_current_backend()
+  if not backend then
+    return nil, nil
+  end
   local lb_alg = get_current_lb_alg()
   local is_sticky = sticky.is_sticky(backend)
 
@@ -164,6 +167,9 @@ function _M.call()
   ngx_balancer.set_more_tries(1)
 
   local host, port = balance()
+  if not host then
+    return error("balancer did not return a host")
+  end
 
   local ok
   ok, err = ngx_balancer.set_current_peer(host, port)

rootfs/etc/nginx/lua/test/balancer_test.lua

@@ -0,0 +1,299 @@
+local cwd = io.popen("pwd"):read('*l')
+package.path = cwd .. "/rootfs/etc/nginx/lua/?.lua;" .. package.path
+
+local balancer, mock_cjson, mock_config, mock_sticky, mock_backends, mock_lrucache, lock, mock_lock,
+  mock_ngx_balancer, mock_ewma
+
+local function dict_generator(vals)
+  local _dict = { __index = {
+      get = function(self, key)
+        return self._vals[key]
+      end,
+      set = function(self, key, val)
+          self._vals[key] = val
+          return true, nil, false
+      end,
+      delete = function(self, key)
+        return self:set(key, nil)
+      end,
+      flush_all = function(self)
+        return
+      end,
+      _vals = vals
+    }
+  }
+  return setmetatable({_vals = vals}, _dict)
+end
+
+local function upstream_generator(name, endpoints, lb_alg)
+  return {
+    name = name,
+    endpoints = endpoints,
+    ["load-balance"] = lb_alg,
+  }
+end
+
+local default_endpoints = {
+  {address = "000.000.000", port = "8080"},
+  {address = "000.000.001", port = "8081"},
+}
+
+local default_upstreams = {
+  mock_rr_upstream = upstream_generator("mock_rr_upstream", default_endpoints, "round_robin"),
+  mock_ewma_upstream = upstream_generator("mock_ewma_upstream", default_endpoints, "ewma"),
+}
+
+local function init()
+  mock_cjson = {}
+  mock_config = {}
+  mock_sticky = {}
+  mock_ngx_balancer = {}
+  mock_ewma = {}
+  mock_backends = dict_generator(default_upstreams)
+  mock_lrucache = {
+    new = function () return mock_backends end
+  }
+  lock = {
+    lock = function() return end,
+    unlock = function() return end
+  }
+  mock_lock = {
+    new = function () return lock end
+  }
+  _G.ngx = {
+    shared = {
+      round_robin_state = dict_generator({}),
+      balancer_ewma = dict_generator({}),
+      balancer_ewma_last_touched_at = dict_generator({}),
+    },
+    var = {},
+    log = function() return end,
+    WARN = "warn",
+    INFO = "info",
+    ERR = "err"
+  }
+  package.loaded["ngx.balancer"] = mock_ngx_balancer
+  package.loaded["resty.lrucache"] = mock_lrucache
+  package.loaded["cjson"] = mock_cjson
+  package.loaded["resty.lock"] = mock_lock
+  package.loaded["balancer.ewma"] = mock_ewma
+  package.loaded["configuration"] = mock_config
+  package.loaded["sticky"] = mock_sticky
+  balancer = require("balancer")
+end
+
+describe("[balancer_test]", function()
+  setup(function()
+    init()
+  end)
+
+  teardown(function()
+    local packages = {"ngx.balancer", "resty.lrucache","cjson", "resty.lock", "balancer.ewma","configuration", "sticky"}
+    for i, package_name in ipairs(packages) do
+      package.loaded[package_name] = nil
+    end
+  end)
+
+  describe("balancer.call():", function()
+    setup(function()
+      mock_ngx_balancer.set_more_tries = function () return end
+      mock_ngx_balancer.set_current_peer = function () return end
+      mock_ewma.after_balance = function () return end
+    end)
+
+    before_each(function()
+      _G.ngx.get_phase = nil
+      _G.ngx.shared.round_robin_state._vals = {}
+      _G.ngx.var = {}
+      mock_backends._vals = default_upstreams
+      mock_sticky.is_sticky = function(b) return false end
+    end)
+
+    describe("phase=log", function()
+      before_each(function()
+        _G.ngx.get_phase = function() return "log" end
+      end)
+
+      it("lb_alg=ewma, ewma_after_balance was called", function()
+        _G.ngx.var.proxy_upstream_name = "mock_ewma_upstream"
+
+        local backend_get_spy = spy.on(mock_backends, "get")
+        local ewma_after_balance_spy = spy.on(mock_ewma, "after_balance")
+
+        assert.has_no_errors(balancer.call)
+        assert.spy(backend_get_spy).was_called_with(match.is_table(), "mock_ewma_upstream")
+        assert.spy(ewma_after_balance_spy).was_called()
+      end)
+
+      it("lb_alg=round_robin, ewma_after_balance was not called", function()
+        _G.ngx.var.proxy_upstream_name = "mock_rr_upstream"
+
+        local backend_get_spy = spy.on(mock_backends, "get")
+        local ewma_after_balance_spy = spy.on(mock_ewma, "after_balance")
+
+        assert.has_no_errors(balancer.call)
+        assert.spy(backend_get_spy).was_called_with(match.is_table(), "mock_rr_upstream")
+        assert.spy(ewma_after_balance_spy).was_not_called()
+      end)
+    end)
+
+    describe("phase=balancer", function()
+      before_each(function()
+        _G.ngx.get_phase = function() return "balancer" end
+      end)
+
+      it("lb_alg=round_robin, peer was successfully set", function()
+        _G.ngx.var.proxy_upstream_name = "mock_rr_upstream"
+
+        local backend_get_spy = spy.on(mock_backends, "get")
+        local set_more_tries_spy = spy.on(mock_ngx_balancer, "set_more_tries")
+        local set_current_peer_spy = spy.on(mock_ngx_balancer, "set_current_peer")
+
+        assert.has_no_errors(balancer.call)
+        assert.spy(backend_get_spy).was_called_with(match.is_table(), "mock_rr_upstream")
+        assert.spy(set_more_tries_spy).was_called_with(1)
+        assert.spy(set_current_peer_spy).was_called_with("000.000.000", "8080")
+
+        mock_backends.get:clear()
+        mock_ngx_balancer.set_more_tries:clear()
+        mock_ngx_balancer.set_current_peer:clear()
+
+        assert.has_no_errors(balancer.call)
+        assert.spy(backend_get_spy).was_called_with(match.is_table(), "mock_rr_upstream")
+        assert.spy(set_more_tries_spy).was_called_with(1)
+        assert.spy(set_current_peer_spy).was_called_with("000.000.001", "8081")
+      end)
+
+      it("lb_alg=ewma, peer was successfully set", function()
+        _G.ngx.var.proxy_upstream_name = "mock_ewma_upstream"
+
+        mock_ewma.balance = function(b) return {address = "000.000.111", port = "8083"} end
+
+        local backend_get_spy = spy.on(mock_backends, "get")
+        local set_more_tries_spy = spy.on(mock_ngx_balancer, "set_more_tries")
+        local set_current_peer_spy = spy.on(mock_ngx_balancer, "set_current_peer")
+
+        assert.has_no_errors(balancer.call)
+        assert.spy(backend_get_spy).was_called_with(match.is_table(), "mock_ewma_upstream")
+        assert.spy(set_more_tries_spy).was_called_with(1)
+        assert.spy(set_current_peer_spy).was_called_with("000.000.111", "8083")
+      end)
+
+      it("sticky=true, returns stored endpoints and peer was successfully set", function()
+        _G.ngx.var.proxy_upstream_name = "mock_rr_upstream"
+
+        mock_sticky.is_sticky = function(b) return true end
+        mock_sticky.get_endpoint = function() return {address = "000.000.011", port = "8082"} end
+
+        local backend_get_spy = spy.on(mock_backends, "get")
+        local set_more_tries_spy = spy.on(mock_ngx_balancer, "set_more_tries")
+        local set_current_peer_spy = spy.on(mock_ngx_balancer, "set_current_peer")
+
+        assert.has_no_errors(balancer.call)
+        assert.spy(backend_get_spy).was_called_with(match.is_table(), "mock_rr_upstream")
+        assert.spy(set_more_tries_spy).was_called_with(1)
+        assert.spy(set_current_peer_spy).was_called_with("000.000.011", "8082")
+      end)
+
+      it("sticky=true, does not return stored endpoints, defaults to round robin", function()
+        _G.ngx.var.proxy_upstream_name = "mock_rr_upstream"
+
+        mock_sticky.is_sticky = function(b) return true end
+        mock_sticky.get_endpoint = function() return nil end
+        mock_sticky.set_endpoint = function(b) return end
+
+        local backend_get_spy = spy.on(mock_backends, "get")
+        local set_more_tries_spy = spy.on(mock_ngx_balancer, "set_more_tries")
+        local set_current_peer_spy = spy.on(mock_ngx_balancer, "set_current_peer")
+
+        assert.has_no_errors(balancer.call)
+        assert.spy(backend_get_spy).was_called_with(match.is_table(), "mock_rr_upstream")
+        assert.spy(set_more_tries_spy).was_called_with(1)
+        assert.spy(set_current_peer_spy).was_called_with("000.000.000", "8080")
+      end)
+
+      it("fails when no backend exists", function()
+        _G.ngx.var.proxy_upstream_name = "mock_rr_upstream"
+
+        mock_backends._vals = {}
+
+        local backend_get_spy = spy.on(mock_backends, "get")
+        local ngx_spy = spy.on(_G.ngx, "log")
+        local set_more_tries_spy = spy.on(mock_ngx_balancer, "set_more_tries")
+        local set_current_peer_spy = spy.on(mock_ngx_balancer, "set_current_peer")
+
+        assert.has_error(balancer.call, "balancer did not return a host")
+        assert.spy(backend_get_spy).was_called_with(match.is_table(), "mock_rr_upstream")
+        assert.spy(set_more_tries_spy).was_called_with(1)
+        assert.spy(set_current_peer_spy).was_not_called()
+      end)
+    end)
+
+    describe("not in phase log or balancer", function()
+      it("returns errors", function()
+        _G.ngx.get_phase = function() return "nope" end
+        assert.has_error(balancer.call, "must be called in balancer or log, but was called in: nope")
+      end)
+    end)
+  end)
+
+  describe("balancer.init_worker():", function()
+    setup(function()
+      _G.ngx.timer = {
+        every = function(interval, func) return func() end
+      }
+      mock_cjson.decode = function(x) return x end
+    end)
+
+    before_each(function()
+      mock_backends._vals = default_upstreams
+    end)
+
+    describe("sync_backends():", function()
+      it("succeeds when no sync is required", function()
+        mock_config.get_backends_data = function() return default_upstreams end
+
+        local backend_set_spy = spy.on(mock_backends, "set")
+
+        assert.has_no_errors(balancer.init_worker)
+        assert.spy(backend_set_spy).was_not_called()
+      end)
+
+      it("lb_alg=round_robin, updates backend when sync is required", function()
+        mock_config.get_backends_data = function() return { default_upstreams.mock_rr_upstream } end
+        mock_backends._vals = {}
+        _G.ngx.shared.round_robin_state._vals = default_upstreams.mock_rr_upstream
+
+        local backend_set_spy = spy.on(mock_backends, "set")
+        local rr_delete_spy = spy.on(_G.ngx.shared.round_robin_state, "delete")
+        local ewma_flush_spy = spy.on(_G.ngx.shared.balancer_ewma, "flush_all")
+        local ewma_lta_flush_spy = spy.on(_G.ngx.shared.balancer_ewma_last_touched_at, "flush_all")
+
+        assert.has_no_errors(balancer.init_worker)
+        assert.spy(backend_set_spy)
+          .was_called_with(match.is_table(), default_upstreams.mock_rr_upstream.name, match.is_table())
+        assert.spy(rr_delete_spy).was_called_with(match.is_table(), default_upstreams.mock_rr_upstream.name)
+        assert.spy(ewma_flush_spy).was_not_called()
+        assert.spy(ewma_lta_flush_spy).was_not_called()
+      end)
+
+      it("lb_alg=ewma, updates backend when sync is required", function()
+        mock_config.get_backends_data = function() return { default_upstreams.mock_ewma_upstream } end
+        mock_backends._vals = {}
+        _G.ngx.shared.round_robin_state._vals = default_upstreams.mock_ewma_upstream
+
+        local backend_set_spy = spy.on(mock_backends, "set")
+        local rr_delete_spy = spy.on(_G.ngx.shared.round_robin_state, "delete")
+        local ewma_flush_spy = spy.on(_G.ngx.shared.balancer_ewma, "flush_all")
+        local ewma_lta_flush_spy = spy.on(_G.ngx.shared.balancer_ewma_last_touched_at, "flush_all")
+
+        assert.has_no_errors(balancer.init_worker)
+        assert.spy(backend_set_spy)
+          .was_called_with(match.is_table(), default_upstreams.mock_ewma_upstream.name, match.is_table())
+        assert.spy(rr_delete_spy).was_called_with(match.is_table(), default_upstreams.mock_ewma_upstream.name)
+        assert.spy(ewma_flush_spy).was_called_with(match.is_table())
+        assert.spy(ewma_lta_flush_spy).was_called_with(match.is_table())
+      end)
+    end)
+  end)
+end)