From b99d7a46a0c2f8fc56e83a767975ea5375e901bf Mon Sep 17 00:00:00 2001 From: Harsh Gupta <42064744+Harshg999@users.noreply.github.com> Date: Fri, 15 Mar 2024 17:37:25 +0530 Subject: [PATCH] DWX-8026: [raz] Improve RAZ HA unit tests (#3651) (cherry picked from commit 4bd2f36b37ef56d5f6d67bb41140739a6793cb13) (cherry picked from commit 87744078b87ab86bf2f1be7d87df75fbe821f55f) Change-Id: Ib97c029298b55a684819e72c2478d2990764499a (cherry picked from commit be1c5fdeacfb5765c0f30c451df62892c481f272) --- .../core/src/desktop/lib/raz/raz_client.py | 2 +- .../src/desktop/lib/raz/raz_client_test.py | 47 ++++++++++++------- 2 files changed, 30 insertions(+), 19 deletions(-) diff --git a/desktop/core/src/desktop/lib/raz/raz_client.py b/desktop/core/src/desktop/lib/raz/raz_client.py index 320dc747c9b..21e24de189c 100644 --- a/desktop/core/src/desktop/lib/raz/raz_client.py +++ b/desktop/core/src/desktop/lib/raz/raz_client.py @@ -171,7 +171,7 @@ def _handle_raz_ha(self, raz_url, headers, data, auth_handler=None, verify=False raz_response = None for r_url in raz_urls_list: - r_url = "%s/api/authz/%s/access?doAs=%s" % (r_url.rstrip('/'), self.service, self.username) + r_url = "%s/api/authz/%s/access?doAs=%s" % (r_url.strip(' ').rstrip('/'), self.service, self.username) LOG.info('Attempting to connect to RAZ URL: %s' % r_url) try: diff --git a/desktop/core/src/desktop/lib/raz/raz_client_test.py b/desktop/core/src/desktop/lib/raz/raz_client_test.py index 31a68da1786..a3c88182dd6 100644 --- a/desktop/core/src/desktop/lib/raz/raz_client_test.py +++ b/desktop/core/src/desktop/lib/raz/raz_client_test.py @@ -14,8 +14,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -import base64 -import sys import unittest from nose.tools import assert_equal, assert_true, assert_raises @@ -23,10 +21,7 @@ from desktop.lib.raz.raz_client import RazClient, get_raz_client 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 +from unittest.mock import patch, Mock class RazClientTest(unittest.TestCase): @@ -34,7 +29,7 @@ class RazClientTest(unittest.TestCase): def setUp(self): self.username = 'gethue' self.raz_url = 'https://raz.gethue.com:8080' - self.raz_urls_ha = 'https://raz_host_1.gethue.com:8080/,https://raz_host_2.gethue.com:8080/' + self.raz_urls_ha = 'https://raz_host_1.gethue.com:8080/, https://raz_host_2.gethue.com:8080/' self.s3_path = 'https://gethue-test.s3.amazonaws.com/gethue/data/customer.csv' self.adls_path = 'https://gethuestorage.dfs.core.windows.net/gethue-container/user/csso_hueuser/customer.csv' @@ -346,12 +341,8 @@ def test_check_access_s3(self): resp = client.check_access(method='GET', url=self.s3_path) - if sys.version_info[0] > 2: - signed_request = 'CiRodHRwczovL2dldGh1ZS10ZXN0LnMzLmFtYXpvbmF3cy5jb20Q' \ - 'ATIYZ2V0aHVlL2RhdGEvY3VzdG9tZXIuY3N2OABCAnMzSgJzMw==' - else: - signed_request = b'CiRodHRwczovL2dldGh1ZS10ZXN0LnMzLmFtYXpvbmF3cy5jb20Q' \ - b'ATIYZ2V0aHVlL2RhdGEvY3VzdG9tZXIuY3N2OABCAnMzSgJzMw==' + signed_request = 'CiRodHRwczovL2dldGh1ZS10ZXN0LnMzLmFtYXpvbmF3cy5jb20Q' \ + 'ATIYZ2V0aHVlL2RhdGEvY3VzdG9tZXIuY3N2OABCAnMzSgJzMw==' requests_post.assert_called_with( 'https://raz.gethue.com:8080/api/authz/s3/access?doAs=gethue', @@ -382,10 +373,10 @@ def test_check_access_s3(self): def test_handle_raz_ha(self): with patch('desktop.lib.sdxaas.knox_jwt.requests_kerberos.HTTPKerberosAuth') as HTTPKerberosAuth: with patch('desktop.lib.raz.raz_client.requests.post') as requests_post: - requests_post.return_value = Mock(status_code=200) request_data = Mock() # Non-HA mode + requests_post.return_value = Mock(status_code=200) client = RazClient(self.raz_url, 'kerberos', username=self.username, service="s3", service_name="cm_s3", cluster_name="cl1") raz_response = client._handle_raz_ha(self.raz_url, auth_handler=HTTPKerberosAuth(), data=request_data, headers={}) @@ -397,9 +388,13 @@ def test_handle_raz_ha(self): verify=False ) assert_equal(raz_response.status_code, 200) + assert_equal(requests_post.call_count, 1) + requests_post.reset_mock() - # HA mode - where first URL sends 200 status code + # HA mode - When RAZ instance1 is healthy and RAZ instance2 is unhealthy client = RazClient(self.raz_urls_ha, 'kerberos', username=self.username, service="s3", service_name="cm_s3", cluster_name="cl1") + + requests_post.side_effect = [Mock(status_code=200), Mock(status_code=404)] raz_response = client._handle_raz_ha(self.raz_urls_ha, auth_handler=HTTPKerberosAuth(), data=request_data, headers={}) requests_post.assert_called_with( @@ -410,12 +405,28 @@ def test_handle_raz_ha(self): verify=False ) assert_equal(raz_response.status_code, 200) + assert_equal(requests_post.call_count, 1) + requests_post.reset_mock() - # When no RAZ URL is healthy - requests_post.return_value = Mock(status_code=404) + # HA mode - When RAZ instance1 is unhealthy and RAZ instance2 is healthy + requests_post.side_effect = [Mock(status_code=404), Mock(status_code=200)] + raz_response = client._handle_raz_ha(self.raz_urls_ha, auth_handler=HTTPKerberosAuth(), data=request_data, headers={}) - client = RazClient(self.raz_urls_ha, 'kerberos', username=self.username, service="s3", service_name="cm_s3", cluster_name="cl1") + requests_post.assert_called_with( + 'https://raz_host_2.gethue.com:8080/api/authz/s3/access?doAs=gethue', + auth=HTTPKerberosAuth(), + headers={}, + json=request_data, + verify=False + ) + assert_equal(raz_response.status_code, 200) + assert_equal(requests_post.call_count, 2) + requests_post.reset_mock() + + # When no RAZ instance is healthy + requests_post.side_effect = [Mock(status_code=404), Mock(status_code=404)] raz_response = client._handle_raz_ha(self.raz_urls_ha, auth_handler=HTTPKerberosAuth(), data=request_data, headers={}) assert_equal(raz_response, None) + assert_equal(requests_post.call_count, 2)