fixes in tests

Author: yrobla

src/codegate/pipeline/pii/pii.py

@@ -81,7 +81,7 @@ def process_results(
             pii_value = text[result.start : result.end]
 
             # add to session store
-            obj = SensitiveData(pii_value, "pii", result.entity_type)
+            obj = SensitiveData(original=pii_value, service="pii", type=result.entity_type)
             uuid_placeholder = self.sensitive_data_manager.store(session_id, obj)
             anonymized_text = anonymized_text.replace(pii_value, uuid_placeholder)
 

src/codegate/pipeline/secrets/secrets.py

@@ -194,7 +194,7 @@ def _hide_secret(self, match: Match) -> str:
         if not match.type:
             raise ValueError("Secret type must be provided")
 
-        obj = SensitiveData(match.value, match.service, match.type)
+        obj = SensitiveData(original=match.value, service=match.service, type=match.type)
         uuid_placeholder = self._sensitive_data_manager.store(self._session_id, obj)
         logger.debug(
             "Stored secret", service=match.service, type=match.type, placeholder=uuid_placeholder

src/codegate/pipeline/sensitive_data/manager.py

@@ -10,10 +10,9 @@
 class SensitiveData(pydantic.BaseModel):
     """Represents sensitive data with additional metadata."""
 
-    def __init__(self, original: str, service: Optional[str], type: Optional[str]):
-        self.original = original
-        self.service = service
-        self.type = type
+    original: str
+    service: Optional[str] = None
+    type: Optional[str] = None
 
 
 class SensitiveDataManager:

tests/pipeline/secrets/test_secrets.py

@@ -186,7 +186,7 @@ def setup_method(self):
     async def test_complete_marker_processing(self):
         """Test processing of a complete REDACTED marker"""
         # Store a secret
-        obj = SensitiveData("secret_value", "test_service", "api_key")
+        obj = SensitiveData(original="secret_value", service="test_service", type="api_key")
         encrypted = self.sensitive_data_manager.store(self.session_id, obj)
 
         # Add content with REDACTED marker to buffer
@@ -272,7 +272,7 @@ async def test_no_markers(self):
     async def test_wrong_session(self):
         """Test unredaction with wrong session ID"""
         # Store secret with one session
-        obj = SensitiveData("test_service", "api_key", "different_session")
+        obj = SensitiveData(original="test_service", service="api_key", type="different_session")
         encrypted = self.sensitive_data_manager.store("different_session", obj)
 
         # Try to unredact with different session

tests/pipeline/sensitive_data/test_manager.py

@@ -22,20 +22,22 @@ def manager(self, mock_session_store):
     def test_store_success(self, manager, mock_session_store):
         """Test storing a SensitiveData object successfully."""
         session_id = "session-123"
-        sensitive_data = SensitiveData("secret_value", "AWS", "API_KEY")
+        sensitive_data = SensitiveData(original="secret_value", service="AWS", type="API_KEY")
 
         # Mock session store behavior
         mock_session_store.add_mapping.return_value = "uuid-123"
 
         result = manager.store(session_id, sensitive_data)
 
         # Verify correct function calls
-        mock_session_store.add_mapping.assert_called_once_with(session_id, sensitive_data.to_json())
+        mock_session_store.add_mapping.assert_called_once_with(
+            session_id, sensitive_data.model_dump_json()
+        )
         assert result == "uuid-123"
 
     def test_store_invalid_session_id(self, manager):
         """Test storing data with an invalid session ID (should return None)."""
-        sensitive_data = SensitiveData("secret_value", "AWS", "API_KEY")
+        sensitive_data = SensitiveData(original="secret_value", service="AWS", type="API_KEY")
         result = manager.store("", sensitive_data)  # Empty session ID
         assert result is None