Enhanced suid checking function to be simpler

Author: luis5tb

src/utils/suid.py

@@ -23,75 +23,31 @@ def check_suid(suid: str) -> bool:
     Returns True if the string is a valid UUID or a llama-stack conversation ID.
 
     Parameters:
-        suid (str | bytes): UUID value to validate — accepts a UUID string,
-        its byte representation, or a llama-stack conversation ID (conv_xxx),
-        or a plain hex string (database format).
+        suid (str): UUID value to validate — accepts a UUID string,
+        or a llama-stack conversation ID (48-char hex, optionally with conv_ prefix).
 
     Notes:
-        Validation is performed by:
-        1. For llama-stack conversation IDs starting with 'conv_':
-           - Strips the 'conv_' prefix
-           - Validates at least 32 hex characters follow (may have additional suffix)
-           - Extracts first 32 hex chars as the UUID part
-           - Converts to UUID format by inserting hyphens at standard positions
-           - Validates the resulting UUID structure
-        2. For plain hex strings (database format, 32+ chars without conv_ prefix):
-           - Validates it's a valid hex string
-           - Extracts first 32 chars as UUID part
-           - Converts to UUID format and validates
-        3. For standard UUIDs: attempts to construct uuid.UUID(suid)
-        Invalid formats or types result in False.
+        Validation accepts:
+        1. Standard UUID format (e.g., '550e8400-e29b-41d4-a716-446655440000')
+        2. 48-character hex string (llama-stack format)
+        3. 'conv_' prefix + 48-character hex string (53 chars total)
     """
-    try:
-        # Accept llama-stack conversation IDs (conv_<hex> format)
-        if isinstance(suid, str) and suid.startswith("conv_"):
-            # Extract the hex string after 'conv_'
-            hex_part = suid[5:]  # Remove 'conv_' prefix
-
-            # Verify it's a valid hex string
-            # llama-stack may use 32 hex chars (UUID) or 36 hex chars (UUID + suffix)
-            if len(hex_part) < 32:
-                return False
-
-            # Verify all characters are valid hex
-            try:
-                int(hex_part, 16)
-            except ValueError:
-                return False
-
-            # Extract the first 32 hex characters (the UUID part)
-            uuid_hex = hex_part[:32]
-
-            # Convert to UUID format with hyphens: 8-4-4-4-12
-            uuid_str = (
-                f"{uuid_hex[:8]}-{uuid_hex[8:12]}-{uuid_hex[12:16]}-"
-                f"{uuid_hex[16:20]}-{uuid_hex[20:]}"
-            )
-
-            # Validate it's a proper UUID
-            uuid.UUID(uuid_str)
+    if not isinstance(suid, str):
+        return False
+
+    # Strip 'conv_' prefix if present
+    hex_part = suid[5:] if suid.startswith("conv_") else suid
+
+    # Check for 48-char hex string (llama-stack conversation ID format)
+    if len(hex_part) == 48:
+        try:
+            int(hex_part, 16)
             return True
+        except ValueError:
+            return False
 
-        # Check if it's a plain hex string (database format without conv_ prefix)
-        if isinstance(suid, str) and len(suid) >= 32:
-            try:
-                int(suid, 16)
-                # Extract the first 32 hex characters (the UUID part)
-                uuid_hex = suid[:32]
-
-                # Convert to UUID format with hyphens: 8-4-4-4-12
-                uuid_str = (
-                    f"{uuid_hex[:8]}-{uuid_hex[8:12]}-{uuid_hex[12:16]}-"
-                    f"{uuid_hex[16:20]}-{uuid_hex[20:]}"
-                )
-
-                # Validate it's a proper UUID
-                uuid.UUID(uuid_str)
-                return True
-            except ValueError:
-                pass  # Not a valid hex string, try standard UUID validation
-
-        # accepts strings and bytes only for UUID validation
+    # Check for standard UUID format
+    try:
         uuid.UUID(suid)
         return True
     except (ValueError, TypeError):

tests/unit/utils/test_suid.py

@@ -1,5 +1,9 @@
 """Unit tests for functions defined in utils.suid module."""
 
+from typing import Any
+
+import pytest
+
 from utils import suid
 
 
@@ -12,16 +16,62 @@ def test_get_suid(self) -> None:
         assert suid.check_suid(suid_value), "Generated SUID is not valid"
         assert isinstance(suid_value, str), "SUID should be a string"
 
-    def test_check_suid_valid(self) -> None:
+    def test_check_suid_valid_uuid(self) -> None:
         """Test that check_suid returns True for a valid UUID."""
         valid_suid = "123e4567-e89b-12d3-a456-426614174000"
+        assert suid.check_suid(valid_suid), "check_suid should return True for UUID"
+
+    def test_check_suid_valid_48char_hex(self) -> None:
+        """Test that check_suid returns True for a 48-char hex string."""
+        valid_hex = "e6afd7aaa97b49ce8f4f96a801b07893d9cb784d72e53e3c"
+        assert len(valid_hex) == 48
         assert suid.check_suid(
-            valid_suid
-        ), "check_suid should return True for a valid SUID"
+            valid_hex
+        ), "check_suid should return True for 48-char hex"
+
+    def test_check_suid_valid_conv_prefix(self) -> None:
+        """Test that check_suid returns True for conv_ + 48-char hex string."""
+        valid_conv = "conv_e6afd7aaa97b49ce8f4f96a801b07893d9cb784d72e53e3c"
+        assert len(valid_conv) == 53
+        assert suid.check_suid(
+            valid_conv
+        ), "check_suid should return True for conv_ prefixed hex"
+
+    def test_check_suid_invalid_string(self) -> None:
+        """Test that check_suid returns False for an invalid string."""
+        assert not suid.check_suid("invalid-uuid")
 
-    def test_check_suid_invalid(self) -> None:
-        """Test that check_suid returns False for an invalid UUID."""
-        invalid_suid = "invalid-uuid"
+    def test_check_suid_valid_32char_hex_uuid(self) -> None:
+        """Test that check_suid returns True for 32-char hex (valid UUID format)."""
+        # 32-char hex is a valid UUID format (without hyphens)
+        assert suid.check_suid("e6afd7aaa97b49ce8f4f96a801b07893")
+
+    def test_check_suid_invalid_hex_wrong_length(self) -> None:
+        """Test that check_suid returns False for hex string with wrong length."""
+        # 47 chars (not 48, not valid UUID)
+        assert not suid.check_suid("e6afd7aaa97b49ce8f4f96a801b07893d9cb784d72e53e3")
+        # 49 chars (not 48, not valid UUID)
+        assert not suid.check_suid("e6afd7aaa97b49ce8f4f96a801b07893d9cb784d72e53e3c1")
+
+    def test_check_suid_invalid_conv_prefix_wrong_length(self) -> None:
+        """Test that check_suid returns False for conv_ with wrong hex length."""
+        # conv_ + 47 chars (not 48)
+        assert not suid.check_suid(
+            "conv_e6afd7aaa97b49ce8f4f96a801b07893d9cb784d72e53e3"
+        )
+        # conv_ + 49 chars (not 48)
         assert not suid.check_suid(
-            invalid_suid
-        ), "check_suid should return False for an invalid SUID"
+            "conv_e6afd7aaa97b49ce8f4f96a801b07893d9cb784d72e53e3c1"
+        )
+
+    def test_check_suid_invalid_non_hex_chars(self) -> None:
+        """Test that check_suid returns False for strings with non-hex characters."""
+        # 48 chars but contains 'g' and 'z'
+        invalid_hex = "g6afd7aaa97b49ce8f4f96a801b07893d9cb784d72e53ezz"
+        assert len(invalid_hex) == 48
+        assert not suid.check_suid(invalid_hex)
+
+    @pytest.mark.parametrize("invalid_type", [None, 123, [], {}])
+    def test_check_suid_invalid_type(self, invalid_type: Any) -> None:
+        """Test that check_suid returns False for non-string types."""
+        assert not suid.check_suid(invalid_type)