[knox_jwt] Improve unit tests for Knox JWT HA support (#3650)

cloudera · Mar 14, 2024 · 08037bd · 08037bd
1 parent 57523a1
commit 08037bd
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 11 deletions.
diff --git a/desktop/core/src/desktop/lib/sdxaas/knox_jwt.py b/desktop/core/src/desktop/lib/sdxaas/knox_jwt.py
@@ -40,6 +40,7 @@ def handle_knox_ha():
     knox_urls_list = knox_urls.split(',')
 
     for k_url in knox_urls_list:
+      k_url = k_url.strip(' ')
       try:
         res = requests.get(k_url.rstrip('/') + _KNOX_TOKEN_API, auth=auth_handler, verify=False)
       except Exception as e:

diff --git a/desktop/core/src/desktop/lib/sdxaas/knox_jwt_test.py b/desktop/core/src/desktop/lib/sdxaas/knox_jwt_test.py
@@ -14,20 +14,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import sys
-import unittest
-
+from unittest.mock import patch, Mock
 from nose.tools import assert_equal, assert_raises
 
 from desktop.conf import SDXAAS
 from desktop.lib.sdxaas.knox_jwt import handle_knox_ha, fetch_jwt
 from desktop.lib.exceptions_renderable import PopupException
 
-if sys.version_info[0] > 2:
-  from unittest.mock import patch, Mock
-else:
-  from mock import patch, Mock
-
 
 def test_handle_knox_ha():
   with patch('desktop.lib.sdxaas.knox_jwt.requests_kerberos.HTTPKerberosAuth') as HTTPKerberosAuth:
@@ -37,33 +30,56 @@ def test_handle_knox_ha():
 
       # Non-HA mode
       reset = SDXAAS.TOKEN_URL.set_for_testing('https://knox-gateway0.gethue.com:8443/dl-name/kt-kerberos/')
-
       try:
         knox_url = handle_knox_ha()
+
         assert_equal(knox_url, 'https://knox-gateway0.gethue.com:8443/dl-name/kt-kerberos/')
+        assert_equal(requests_get.call_count, 0) # Simply returning the URL string
       finally:
         reset()
+        requests_get.reset_mock()
 
-      # HA mode - where first URL sends 200 status code
+      # HA mode - When gateway0 is healthy and gateway1 is unhealthy
+      requests_get.side_effect = [Mock(status_code=200), Mock(status_code=404)]
       reset = SDXAAS.TOKEN_URL.set_for_testing(
         'https://knox-gateway0.gethue.com:8443/dl-name/kt-kerberos/, https://knox-gateway1.gethue.com:8443/dl-name/kt-kerberos/')
 
       try:
         knox_url = handle_knox_ha()
+
         assert_equal(knox_url, 'https://knox-gateway0.gethue.com:8443/dl-name/kt-kerberos/')
+        assert_equal(requests_get.call_count, 1)
       finally:
         reset()
+        requests_get.reset_mock()
+
+      # HA mode - When gateway0 is unhealthy and gateway1 is healthy
+      requests_get.side_effect = [Mock(status_code=404), Mock(status_code=200)]
+      reset = SDXAAS.TOKEN_URL.set_for_testing(
+        'https://knox-gateway0.gethue.com:8443/dl-name/kt-kerberos/, https://knox-gateway1.gethue.com:8443/dl-name/kt-kerberos/')
 
-      # When no Knox URL is healthy
+      try:
+        knox_url = handle_knox_ha()
+
+        assert_equal(knox_url, 'https://knox-gateway1.gethue.com:8443/dl-name/kt-kerberos/')
+        assert_equal(requests_get.call_count, 2)
+      finally:
+        reset()
+        requests_get.reset_mock()
+
+      # When both gateway0 and gateway1 are unhealthy
       requests_get.return_value = Mock(status_code=404)
       reset = SDXAAS.TOKEN_URL.set_for_testing(
         'https://knox-gateway0.gethue.com:8443/dl-name/kt-kerberos/, https://knox-gateway1.gethue.com:8443/dl-name/kt-kerberos/')
 
       try:
         knox_url = handle_knox_ha()
+
         assert_equal(knox_url, None)
+        assert_equal(requests_get.call_count, 2)
       finally:
         reset()
+        requests_get.reset_mock()
 
 
 def test_fetch_jwt():