Add unit test cases for balancer lua module

rootfs/etc/nginx/lua/balancer.lua

@@ -257,6 +257,7 @@ end
 if _TEST then
   _M.get_implementation = get_implementation
   _M.sync_backend = sync_backend
+  _M.route_to_alternative_balancer = route_to_alternative_balancer
 end
 
 return _M

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

@@ -1,6 +1,17 @@
 _G._TEST = true
 
 local balancer, expected_implementations, backends
+local original_ngx = ngx
+
+local function reset_ngx()
+  _G.ngx = original_ngx
+end
+
+local function mock_ngx(mock)
+  local _ngx = mock
+  setmetatable(_ngx, { __index = ngx })
+  _G.ngx = _ngx
+end
 
 local function reset_balancer()
   package.loaded["balancer"] = nil
@@ -30,6 +41,12 @@ local function reset_backends()
         { address = "10.184.98.239", port = "8080", maxFails = 0, failTimeout = 0 },
       },
       sessionAffinityConfig = { name = "", cookieSessionAffinity = { name = "" } },
+      trafficShapingPolicy = {
+        weight = 0,
+        header = "",
+        headerValue = "",
+        cookie = ""
+      },
     },
     { name = "my-dummy-app-1", ["load-balance"] = "round_robin", },
     { 
@@ -55,6 +72,10 @@ describe("Balancer", function()
     reset_backends()
   end)
 
+  after_each(function()
+    reset_ngx()
+  end)
+
   describe("get_implementation()", function()
     it("returns correct implementation for given backend", function()
       for _, backend in pairs(backends) do
@@ -65,6 +86,174 @@ describe("Balancer", function()
     end)
   end)
 
+  describe("route_to_alternative_balancer()", function()
+    local backend, _balancer
+
+    before_each(function()
+      backend = backends[1]
+      _balancer = {
+        alternative_backends = {
+          backend.name,
+        }
+      }
+      mock_ngx({ var = { request_uri = "/" } })
+    end)
+
+    it("returns false when no trafficShapingPolicy is set", function()
+      balancer.sync_backend(backend)
+      assert.equal(false, balancer.route_to_alternative_balancer(_balancer))
+    end)
+
+    it("returns false when no alternative backends is set", function()
+      backend.trafficShapingPolicy.weight = 100
+      balancer.sync_backend(backend)
+      _balancer.alternative_backends = nil
+      assert.equal(false, balancer.route_to_alternative_balancer(_balancer))
+    end)
+
+    it("returns false when alternative backends name does not match", function()
+      backend.trafficShapingPolicy.weight = 100
+      balancer.sync_backend(backend)
+      _balancer.alternative_backends[1] = "nonExistingBackend"
+      assert.equal(false, balancer.route_to_alternative_balancer(_balancer))
+    end)
+
+    context("canary by weight", function()
+      it("returns true when weight is 100", function()
+        backend.trafficShapingPolicy.weight = 100
+        balancer.sync_backend(backend)
+        assert.equal(true, balancer.route_to_alternative_balancer(_balancer))
+      end)
+
+      it("returns false when weight is 0", function()
+        backend.trafficShapingPolicy.weight = 0
+        balancer.sync_backend(backend)
+        assert.equal(false, balancer.route_to_alternative_balancer(_balancer))
+      end)
+    end)
+
+    context("canary by cookie", function()
+      it("returns correct result for given cookies", function()
+        backend.trafficShapingPolicy.cookie = "canaryCookie"
+        balancer.sync_backend(backend)
+        local test_patterns = {
+          {
+            case_title = "cookie_value is 'always'",
+            request_cookie_name = "canaryCookie",
+            request_cookie_value = "always",
+            expected_result = true,
+          },
+          {
+            case_title = "cookie_value is 'never'",
+            request_cookie_name = "canaryCookie",
+            request_cookie_value = "never",
+            expected_result = false,
+          },
+          {
+            case_title = "cookie_value is undefined",
+            request_cookie_name = "canaryCookie",
+            request_cookie_value = "foo",
+            expected_result = false,
+          },
+          {
+            case_title = "cookie_name is undefined",
+            request_cookie_name = "foo",
+            request_cookie_value = "always",
+            expected_result = false
+          },
+        }
+        for _, test_pattern in pairs(test_patterns) do
+          mock_ngx({ var = {
+            ["cookie_" .. test_pattern.request_cookie_name] = test_pattern.request_cookie_value,
+            request_uri = "/"
+          }})
+          assert.message("\nTest data pattern: " .. test_pattern.case_title)
+            .equal(test_pattern.expected_result, balancer.route_to_alternative_balancer(_balancer))
+          reset_ngx()
+        end
+      end)
+    end)
+
+    context("canary by header", function()
+      it("returns correct result for given headers", function()
+        local test_patterns = {
+          -- with no header value setting
+          {
+            case_title = "no custom header value and header value is 'always'",
+            header_name = "canaryHeader",
+            header_value = "",
+            request_header_name = "canaryHeader",
+            request_header_value = "always",
+            expected_result = true,
+          },
+          {
+            case_title = "no custom header value and header value is 'never'",
+            header_name = "canaryHeader",
+            header_value = "",
+            request_header_name = "canaryHeader",
+            request_header_value = "never",
+            expected_result = false,
+          },
+          {
+            case_title = "no custom header value and header value is undefined",
+            header_name = "canaryHeader",
+            header_value = "",
+            request_header_name = "canaryHeader",
+            request_header_value = "foo",
+            expected_result = false,
+          },
+          {
+            case_title = "no custom header value and header name is undefined",
+            header_name = "canaryHeader",
+            header_value = "",
+            request_header_name = "foo",
+            request_header_value = "always",
+            expected_result = false,
+          },
+          -- with header value setting
+          {
+            case_title = "custom header value is set and header value is 'always'",
+            header_name = "canaryHeader",
+            header_value = "foo",
+            request_header_name = "canaryHeader",
+            request_header_value = "always",
+            expected_result = false,
+          },
+          {
+            case_title = "custom header value is set and header value match custom header value",
+            header_name = "canaryHeader",
+            header_value = "foo",
+            request_header_name = "canaryHeader",
+            request_header_value = "foo",
+            expected_result = true,
+          },
+          {
+            case_title = "custom header value is set and header name is undefined",
+            header_name = "canaryHeader",
+            header_value = "foo",
+            request_header_name = "bar",
+            request_header_value = "foo",
+            expected_result = false
+          },
+        }
+
+        for _, test_pattern in pairs(test_patterns) do
+          reset_balancer()
+          backend.trafficShapingPolicy.header = test_pattern.header_name
+          backend.trafficShapingPolicy.headerValue = test_pattern.header_value
+          balancer.sync_backend(backend)
+          mock_ngx({ var = {
+            ["http_" .. test_pattern.request_header_name] = test_pattern.request_header_value,
+            request_uri = "/"
+          }})
+          assert.message("\nTest data pattern: " .. test_pattern.case_title)
+            .equal(test_pattern.expected_result, balancer.route_to_alternative_balancer(_balancer))
+          reset_ngx()
+        end
+      end)
+    end)
+  end)
+
   describe("sync_backend()", function()
     local backend, implementation