SK-3039:Update context to other valid data types.

Author: yaswanth-pula-skyflow

skyflow/utils/validations/_validations.py

@@ -115,8 +115,10 @@ def validate_token_options(logger, credentials, config_id_type=None, config_id=N
             SkyflowMessages.Error.EMPTY_CONTEXT_IN_CONFIG.value.format(config_id_type, config_id)
             if config_id_type and config_id else SkyflowMessages.Error.EMPTY_CONTEXT.value
         )
+        # Scalars are accepted because the token engine and the auth service both support
+        # them as ctx claims - bool is listed explicitly even though it is an int subclass.
         validate_required_field(
-            logger, credentials, CredentialField.CONTEXT, (str, dict),
+            logger, credentials, CredentialField.CONTEXT, (str, dict, bool, int, float),
             empty_context_error,
             SkyflowMessages.Error.INVALID_CONTEXT_IN_CONFIG.value.format(config_id_type, config_id)
             if config_id_type and config_id else SkyflowMessages.Error.INVALID_CONTEXT.value

tests/utils/validations/test__validations.py

@@ -186,15 +186,31 @@ def test_validate_credentials_with_invalid_dict_context_key(self):
             SkyflowMessages.Error.INVALID_CTX_MAP_KEY.value.format("invalid key")
         )
 
-    def test_validate_credentials_with_invalid_context_type(self):
-        for invalid_context in [123, ["user_12345"], True]:
-            credentials = {
-                "api_key": "sky-abc12-1234567890abcdef1234567890abcdef",
-                "context": invalid_context
-            }
-            with self.assertRaises(SkyflowError) as context:
+    def test_validate_credentials_with_scalar_context(self):
+        """Numbers and booleans are valid ctx claims, so they must not be rejected here.
+
+        The token engine and the auth service both accept them; rejecting them would make the
+        client stricter than generate_bearer_token for no reason. Falsy scalars are included
+        because nothing on this path may test truthiness.
+        """
+        for scalar_context in [123, 0, 1.5, 0.0, True, False]:
+            with self.subTest(context=scalar_context):
+                credentials = {
+                    "api_key": "sky-abc12-1234567890abcdef1234567890abcdef",
+                    "context": scalar_context
+                }
                 validate_credentials(self.logger, credentials)
-            self.assertEqual(context.exception.message, SkyflowMessages.Error.INVALID_CONTEXT.value)
+
+    def test_validate_credentials_with_invalid_context_type(self):
+        for invalid_context in [["user_12345"], ("user_12345",), None, object()]:
+            with self.subTest(context=invalid_context):
+                credentials = {
+                    "api_key": "sky-abc12-1234567890abcdef1234567890abcdef",
+                    "context": invalid_context
+                }
+                with self.assertRaises(SkyflowError) as context:
+                    validate_credentials(self.logger, credentials)
+                self.assertEqual(context.exception.message, SkyflowMessages.Error.INVALID_CONTEXT.value)
 
     def test_validate_credentials_with_dict_context_in_config(self):
         """Config-scoped messages are used when a config id is available."""

tests/vault/client/test__client.py

@@ -345,6 +345,27 @@ def test_get_bearer_token_forwards_dict_context_unmodified(self, mock_generate):
         self.assertEqual(options["ctx"], DICT_CONTEXT)
         self.assertNotIn("role_ids", options)
 
+    @patch("skyflow.vault.client.client.generate_bearer_token", return_value=("sa_token", None))
+    def test_get_bearer_token_forwards_scalar_context_unmodified(self, mock_generate):
+        """Numbers and booleans reach the token engine as-is, with their type preserved.
+
+        A ctx of 123 and a ctx of "123" are different JWT claims, so a coercion here would
+        silently change which policies match. Falsy scalars are covered too - they must not
+        be dropped by a truthiness check.
+        """
+        for scalar_context in [123, 0, 1.5, 0.0, True, False]:
+            with self.subTest(context=scalar_context):
+                mock_generate.reset_mock()
+                self.vault_client._VaultClient__bearer_token = None
+                credentials = {**CREDENTIALS_WITH_PATH, "context": scalar_context}
+
+                self.vault_client.get_bearer_token(credentials)
+
+                _, options, _ = mock_generate.call_args[0]
+                self.assertEqual(options["ctx"], scalar_context)
+                # assertEqual alone would pass on True vs 1, so pin the exact type.
+                self.assertIs(type(options["ctx"]), type(scalar_context))
+
     @patch("skyflow.vault.client.client.generate_bearer_token", return_value=("sa_token", None))
     def test_get_bearer_token_omits_options_when_roles_and_context_absent(self, mock_generate):
         self.vault_client.get_bearer_token(CREDENTIALS_WITH_PATH)