diff --git a/lib/model/generic_transformer.py b/lib/model/generic_transformer.py index c8f3fc8..04c0a5e 100644 --- a/lib/model/generic_transformer.py +++ b/lib/model/generic_transformer.py @@ -76,9 +76,3 @@ def vectorize(self, texts: List[str]) -> List[List[float]]: """ return self.model.encode(texts).tolist() - def handle_fingerprinting_error(self, error: Exception): - """ - Handle any error that occurs during vectorization. - """ - logger.error(f"Error during vectorization: {error}") - raise error diff --git a/test/lib/model/test_generic.py b/test/lib/model/test_generic.py index 1ad3f21..262b373 100644 --- a/test/lib/model/test_generic.py +++ b/test/lib/model/test_generic.py @@ -1,3 +1,4 @@ +import traceback import os import unittest from unittest.mock import MagicMock, patch @@ -108,13 +109,11 @@ def test_separate_cached_docs(self): self.assertEqual(texts_to_vectorize[0], "Hello") self.assertEqual(docs[1].body.result, [4, 5, 6]) - @patch('lib.model.generic_transformer.logger') - def test_handle_fingerprinting_error(self, mock_logger): - with self.assertRaises(Exception) as context: - self.model.handle_fingerprinting_error(ValueError("An error occurred")) - - mock_logger.error.assert_called_once_with("Error during vectorization: An error occurred") - self.assertTrue(isinstance(context.exception, ValueError)) + @patch('lib.sentry.capture_custom_message') + def test_handle_fingerprinting_error(self, mock_capture_custom_message): + error = ValueError("An error occurred") + response = self.model.handle_fingerprinting_error(error) + self.assertIsInstance(response, schemas.ErrorResponse) if __name__ == '__main__': unittest.main()