SK-2813: add unit tests

Author: saileshwar-skyflow

tests/service_account/test__utils.py

@@ -57,6 +57,9 @@ class TestServiceAccountUtils(unittest.TestCase):
 
     # ── is_expired ────────────────────────────────────────────────────────────
 
+    def test_is_expired_none_token(self):
+        self.assertTrue(is_expired(None))
+
     def test_is_expired_empty_token(self):
         self.assertTrue(is_expired(""))
 
@@ -160,6 +163,18 @@ def test_get_service_account_token_invalid_token_uri(self):
             get_service_account_token(creds, {}, None)
         self.assertEqual(context.exception.message, SkyflowMessages.Error.INVALID_TOKEN_URI.value)
 
+    def test_get_service_account_token_invalid_token_uri_in_options(self):
+        creds = {
+            'privateKey': 'key',
+            'clientID': 'id',
+            'keyID': 'kid',
+            'tokenURI': 'https://valid-url.com',
+        }
+        options = {'token_uri': 'not-a-valid-url'}
+        with self.assertRaises(SkyflowError) as context:
+            get_service_account_token(creds, options, None)
+        self.assertEqual(context.exception.message, SkyflowMessages.Error.INVALID_TOKEN_URI.value)
+
     @patch("skyflow.service_account._utils.AuthClient")
     @patch("skyflow.service_account._utils.get_signed_jwt")
     def test_get_service_account_token_with_role_ids_formats_scope(self, mock_get_signed_jwt, mock_auth_client):

tests/utils/test__utils.py

@@ -1,7 +1,6 @@
 import unittest
-from unittest.mock import patch, Mock
+from unittest.mock import patch, Mock, MagicMock, PropertyMock
 import os
-from unittest.mock import MagicMock
 from urllib.parse import quote
 import tempfile, json
 from requests import PreparedRequest
@@ -15,7 +14,7 @@
     parse_detokenize_response, parse_tokenize_response, parse_query_response, parse_invoke_connection_response, \
     handle_exception, validate_api_key, encode_column_values, parse_deidentify_text_response, \
     parse_reidentify_text_response, convert_detected_entity_to_entity_info
-from skyflow.utils._utils import parse_path_params, to_lowercase_keys, get_metrics, handle_json_error
+from skyflow.utils._utils import parse_path_params, to_lowercase_keys, get_metrics, handle_json_error, r_urlencode
 from skyflow.utils.enums import EnvUrls, Env, ContentType
 from skyflow.vault.connection import InvokeConnectionResponse
 from skyflow.vault.data import InsertResponse, DeleteResponse, GetResponse, QueryResponse
@@ -36,6 +35,13 @@ def test_get_credentials_env_variable(self):
         credentials_string = credentials.get('credentials_string')
         self.assertEqual(credentials_string, json.dumps(VALID_ENV_CREDENTIALS).replace('\n', '\\n'))
 
+    @patch("skyflow.utils._utils.dotenv.find_dotenv", return_value=None)
+    @patch.dict(os.environ, {}, clear=True)
+    def test_get_credentials_no_credentials_raises(self, mock_find_dotenv):
+        with self.assertRaises(SkyflowError) as context:
+            get_credentials(config_level_creds=None, common_skyflow_creds=None)
+        self.assertEqual(context.exception.message, SkyflowMessages.Error.INVALID_CREDENTIALS.value)
+
     def test_get_credentials_with_config_level_creds(self):
         test_creds = {"authToken": "test_token"}
         creds = get_credentials(config_level_creds=test_creds)
@@ -140,13 +146,54 @@ def test_to_lowercase_keys(self):
         expected_output = {"key1": "value1", "key2": "value2"}
         self.assertEqual(to_lowercase_keys(input_dict), expected_output)
 
+    def test_r_urlencode_with_list_input(self):
+        pairs = {}
+        r_urlencode([], pairs, ["a", "b"])
+        self.assertIn("[0]", pairs)
+        self.assertIn("[1]", pairs)
+        self.assertEqual(pairs["[0]"], "a")
+        self.assertEqual(pairs["[1]"], "b")
+
+    def test_r_urlencode_with_tuple_input(self):
+        pairs = {}
+        r_urlencode([], pairs, ("x", "y"))
+        self.assertIn("[0]", pairs)
+        self.assertEqual(pairs["[0]"], "x")
+
     def test_get_metrics(self):
         metrics = get_metrics()
         self.assertIn('sdk_name_version', metrics)
         self.assertIn('sdk_client_device_model', metrics)
         self.assertIn('sdk_client_os_details', metrics)
         self.assertIn('sdk_runtime_details', metrics)
 
+    def test_get_metrics_platform_node_exception(self):
+        import skyflow.utils._utils as utils_module
+        utils_module._CACHED_METRICS.clear()
+        with patch("skyflow.utils._utils.platform") as mock_platform:
+            mock_platform.node.side_effect = OSError("no node")
+            metrics = utils_module.get_metrics()
+        self.assertEqual(metrics["sdk_client_device_model"], "")
+        utils_module._CACHED_METRICS.clear()
+
+    def test_get_metrics_sys_attribute_exception(self):
+        import skyflow.utils._utils as utils_module
+        utils_module._CACHED_METRICS.clear()
+
+        class _RaisingSys:
+            @property
+            def platform(self):
+                raise RuntimeError("no platform")
+            @property
+            def version(self):
+                raise RuntimeError("no version")
+
+        with patch("skyflow.utils._utils.sys", _RaisingSys()):
+            metrics = utils_module.get_metrics()
+        self.assertEqual(metrics["sdk_client_os_details"], "")
+        self.assertIn("sdk_runtime_details", metrics)
+        utils_module._CACHED_METRICS.clear()
+
 
     def test_construct_invoke_connection_request_valid(self):
         mock_connection_request = Mock()
@@ -244,6 +291,16 @@ def test_construct_invoke_connection_request_with_form_date_content_type(self):
 
         self.assertIsInstance(result, PreparedRequest)
 
+    def test_parse_insert_response_with_tokens_continue_on_error(self):
+        api_response = Mock()
+        api_response.headers = {"x-request-id": "req-1"}
+        api_response.data = Mock(responses=[
+            {"Status": 200, "Body": {"records": [{"skyflow_id": "id1", "tokens": {"col1": "tok1"}}]}},
+        ])
+        result = parse_insert_response(api_response, continue_on_error=True)
+        self.assertEqual(result.inserted_fields[0]["col1"], "tok1")
+        self.assertEqual(result.inserted_fields[0]["skyflow_id"], "id1")
+
     def test_parse_insert_response(self):
         api_response = Mock()
         api_response.headers = {"x-request-id": "12345", "content-type": "application/json"}
@@ -423,6 +480,31 @@ def test_parse_invoke_connection_response_http_error_with_json_error_message(sel
         self.assertEqual(context.exception.message, "Not Found")
         self.assertEqual(context.exception.request_id, "1234")
 
+    @patch("requests.Response")
+    def test_parse_invoke_connection_response_with_error_from_client_header(self, mock_response):
+        from requests.models import HTTPError
+        mock_response.status_code = 400
+        mock_response.content = json.dumps({
+            "error": {
+                "message": "Client error",
+                "http_code": 400,
+                "http_status": "Bad Request",
+                "grpc_code": 3,
+                "details": None,
+            }
+        }).encode("utf-8")
+        mock_response.headers = {
+            "x-request-id": "rid-1",
+            "error-from-client": "true",
+        }
+        mock_response.raise_for_status.side_effect = HTTPError("400")
+        with self.assertRaises(SkyflowError) as context:
+            parse_invoke_connection_response(mock_response)
+        err = context.exception
+        self.assertEqual(err.message, "Client error")
+        self.assertIsNotNone(err.details)
+        self.assertTrue(any(d.get("error_from_client") is True for d in err.details))
+
     @patch("requests.Response")
     def test_parse_invoke_connection_response_http_error_without_json_error_message(self, mock_response):
         mock_response.status_code = 500

tests/vault/controller/test__detect.py

@@ -2,6 +2,7 @@
 from unittest.mock import Mock, patch, MagicMock
 import base64
 import os
+import tempfile
 from skyflow.error import SkyflowError
 from skyflow.generated.rest import WordCharacterCount
 from skyflow.utils import SkyflowMessages
@@ -707,3 +708,98 @@ def test_deidentify_file_using_file_path(self, mock_open, mock_basename, mock_ba
             self.assertIsNone(result.page_count)
             self.assertIsNone(result.slide_count)
             self.assertEqual(result.entities, [])
+
+    def test_poll_for_processed_file_exception(self):
+        files_api = Mock()
+        files_api.with_raw_response = files_api
+        files_api.get_run.side_effect = Exception("poll error")
+        self.vault_client.get_detect_file_api.return_value = files_api
+        with self.assertRaises(Exception):
+            self.detect._Detect__poll_for_processed_file("runid", max_wait_time=5)
+
+    def test_save_output_directory_not_exists(self):
+        output = Mock()
+        output.processedFile = base64.b64encode(b"data").decode()
+        output.processedFileType = "redacted_file"
+        output.processedFileExtension = "txt"
+        response = Mock()
+        response.output = [output]
+        with patch("skyflow.vault.controller._detect.os.path.exists", return_value=False):
+            self.detect._Detect__save_deidentify_file_response_output(
+                response, "/nonexistent_dir", "file.txt", "file"
+            )
+
+    def test_save_output_second_non_redacted_item(self):
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            output1 = Mock()
+            output1.processedFile = base64.b64encode(b"data1").decode()
+            output1.processedFileType = "redacted_file"
+            output1.processedFileExtension = "txt"
+            output2 = Mock()
+            output2.processedFile = base64.b64encode(b"data2").decode()
+            output2.processedFileType = "entities"
+            output2.processedFileExtension = "json"
+            response = Mock()
+            response.output = [output1, output2]
+            self.detect._Detect__save_deidentify_file_response_output(
+                response, tmp_dir, "original.txt", "original"
+            )
+
+    def test_save_output_path_traversal_blocked(self):
+        output = Mock()
+        output.processedFile = base64.b64encode(b"data").decode()
+        output.processedFileType = "redacted_file"
+        output.processedFileExtension = "txt"
+        response = Mock()
+        response.output = [output]
+        call_count = [0]
+
+        def fake_realpath(p):
+            call_count[0] += 1
+            if call_count[0] == 1:
+                return "/safe_dir"
+            return "/outside/path"
+
+        with patch("skyflow.vault.controller._detect.os.path.exists", return_value=True), \
+             patch("skyflow.vault.controller._detect.os.path.realpath", side_effect=fake_realpath):
+            self.detect._Detect__save_deidentify_file_response_output(
+                response, "/safe_dir", "file.txt", "file"
+            )
+
+    def test_save_output_write_exception(self):
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            output = Mock()
+            output.processedFile = base64.b64encode(b"data").decode()
+            output.processedFileType = "redacted_file"
+            output.processedFileExtension = "txt"
+            response = Mock()
+            response.output = [output]
+            with patch("skyflow.vault.controller._detect.base64.b64decode",
+                       side_effect=Exception("decode error")), \
+                 self.assertRaises(Exception):
+                self.detect._Detect__save_deidentify_file_response_output(
+                    response, tmp_dir, "file.txt", "file"
+                )
+
+    @patch("skyflow.vault.controller._detect.validate_deidentify_file_request")
+    @patch("skyflow.vault.controller._detect.base64")
+    def test_deidentify_file_api_error_inside_try(self, mock_base64, mock_validate):
+        file_content = b"test content"
+        file_obj = Mock()
+        file_obj.read.return_value = file_content
+        file_obj.name = "test.txt"
+        mock_base64.b64encode.return_value.decode.return_value = "encoded"
+        req = DeidentifyFileRequest(file=FileInput(file=file_obj))
+        req.entities = []
+        req.token_format = None
+        req.allow_regex_list = []
+        req.restrict_regex_list = []
+        req.transformations = None
+        req.output_directory = None
+        req.wait_time = None
+        files_api = Mock()
+        files_api.with_raw_response = files_api
+        files_api.deidentify_text.side_effect = Exception("API error inside try")
+        self.vault_client.get_detect_file_api.return_value = files_api
+        with self.assertRaises(Exception):
+            self.detect.deidentify_file(req)