redo test fixtures and readd specific schemas

lib/http.py

@@ -35,7 +35,7 @@ def process_item(process_name: str, message: Dict[str, Any]):
     logger.info(message)
     queue_prefix = Queue.get_queue_prefix()
     queue = QueueWorker.create(process_name)
-    queue.push_message(f"{queue_prefix}{process_name}", schemas.Message(body=message))
+    queue.push_message(f"{queue_prefix}{process_name}", schemas.Message(body=message, model_name=process_name))
     return {"message": "Message pushed successfully", "queue": process_name, "body": message}
 
 @app.post("/trigger_callback")

lib/schemas.py

@@ -3,13 +3,13 @@
 
 # Output hash values can be of different types.
 class GenericItem(BaseModel):
-    id: Optional[str] = None
+    id: str
     callback_url: Optional[str] = None
     url: Optional[str] = None
     text: Optional[str] = None
     raw: Optional[dict] = {}
 
-class MediaItem(BaseModel):
+class MediaItem(GenericItem):
     hash_value: Optional[Any] = None
 
 class VideoItem(MediaItem):
@@ -24,7 +24,7 @@ def set_body(cls, values):
         body = values.get("body")
         model_name = values.get("model_name")
         if model_name == "video__Model":
-            values["body"] = VideoItem(**values["body"])
+            values["body"] = VideoItem(**values["body"]).dict()
         if model_name in ["audio__Model", "image__Model", "fptg__Model", "indian_sbert__Model", "mean_tokens__Model", "fasttext__Model"]:
-            values["body"] = MediaItem(**values["body"])
+            values["body"] = MediaItem(**values["body"]).dict()
         return values

test/lib/model/test_audio.py

@@ -24,7 +24,7 @@ def test_process_audio_success(self, mock_fingerprint_file, mock_request, mock_u
 
         mock_urlopen.return_value = MagicMock(read=MagicMock(return_value=contents))
 
-        audio = schemas.Message(body=schemas.GenericItem(id="123", callback_url="http://example.com/callback", url="https://example.com/audio.mp3"))
+        audio = schemas.Message(body={"id": "123", "callback_url": "http://example.com/callback", "url": "https://example.com/audio.mp3"}, model_name="audio__Model")
         result = self.audio_model.process(audio)
         mock_request.assert_called_once_with(audio.body.url, headers={'User-Agent': 'Mozilla/5.0'})
         mock_urlopen.assert_called_once_with(mock_request)
@@ -45,7 +45,7 @@ def test_process_audio_failure(self, mock_decode_fingerprint, mock_fingerprint_f
 
         mock_urlopen.return_value = MagicMock(read=MagicMock(return_value=contents))
 
-        audio = schemas.Message(body=schemas.GenericItem(id="123", callback_url="http://example.com/callback", url="https://example.com/audio.mp3"))
+        audio = schemas.Message(body={"id": "123", "callback_url": "http://example.com/callback", "url": "https://example.com/audio.mp3"}, model_name="audio__Model")
         result = self.audio_model.process(audio)
         mock_request.assert_called_once_with(audio.body.url, headers={'User-Agent': 'Mozilla/5.0'})
         mock_urlopen.assert_called_once_with(mock_request)

test/lib/model/test_fasttext.py

@@ -15,7 +15,7 @@ def test_respond(self, mock_fasttext_load_model, mock_hf_hub_download):
         mock_fasttext_load_model.return_value = self.mock_model
         self.mock_model.predict.return_value = (['__label__eng_Latn'], np.array([0.9]))
         model = FasttextModel()  # Now it uses mocked functions
-        query = [schemas.Message(body=schemas.GenericItem(id="123", callback_url="http://example.com/callback", text="Hello, how are you?"))]
+        query = [schemas.Message(body={"id": "123", "callback_url": "http://example.com/callback", "text": "Hello, how are you?"}, model_name="fasttext__Model")]
         response = model.respond(query)
         self.assertEqual(len(response), 1)
         self.assertEqual(response[0].body.hash_value, {'language': 'en', 'script': None, 'score': 0.9})

test/lib/model/test_fptg.py

@@ -13,7 +13,7 @@ def setUp(self):
         self.mock_model = MagicMock()
 
     def test_vectorize(self):
-        texts = [schemas.Message(body=schemas.GenericItem(id="123", callback_url="http://example.com/callback", text="Hello, how are you?")), schemas.Message(body=schemas.GenericItem(id="123", callback_url="http://example.com/callback", text="I'm doing great, thanks!"))]
+        texts = [schemas.Message(body={"id": "123", "callback_url": "http://example.com/callback", "text": "Hello, how are you?"}, model_name="fptg__Model"), schemas.Message(body={"id": "123", "callback_url": "http://example.com/callback", "text": "I'm doing great, thanks!"}, model_name="fptg__Model")]
         self.model.model = self.mock_model
         self.model.model.encode = MagicMock(return_value=np.array([[4, 5, 6], [7, 8, 9]]))
         vectors = self.model.vectorize(texts)
@@ -22,7 +22,7 @@ def test_vectorize(self):
         self.assertEqual(vectors[1], [7, 8, 9])
 
     def test_respond(self):
-        query = schemas.Message(body=schemas.GenericItem(id="123", callback_url="http://example.com/callback", text="Anong pangalan mo?"))
+        query = schemas.Message(body={"id": "123", "callback_url": "http://example.com/callback", "text": "Anong pangalan mo?"}, model_name="fptg__Model")
         self.model.vectorize = MagicMock(return_value=[[1, 2, 3]])
         response = self.model.respond(query)
         self.assertEqual(len(response), 1)

test/lib/model/test_generic.py

@@ -13,7 +13,7 @@ def setUp(self):
         self.mock_model = MagicMock()
 
     def test_vectorize(self):
-        texts = [schemas.Message(body=schemas.GenericItem(id="123", callback_url="http://example.com/callback", text="Hello, how are you?")), schemas.Message(body=schemas.GenericItem(id="123", callback_url="http://example.com/callback", text="I'm doing great, thanks!"))]
+        texts = [schemas.Message(body={"id": "123", "callback_url": "http://example.com/callback", "text": "Hello, how are you?"}, model_name="fptg__Model"), schemas.Message(body={"id": "123", "callback_url": "http://example.com/callback", "text": "I'm doing great, thanks!"}, model_name="fptg__Model")]
         self.model.model = self.mock_model
         self.model.model.encode = MagicMock(return_value=np.array([[4, 5, 6], [7, 8, 9]]))
         vectors = self.model.vectorize(texts)
@@ -22,7 +22,7 @@ def test_vectorize(self):
         self.assertEqual(vectors[1], [7, 8, 9])
 
     def test_respond(self):
-        query = schemas.Message(body=schemas.GenericItem(id="123", callback_url="http://example.com/callback", text="Anong pangalan mo?"))
+        query = schemas.Message(body={"id": "123", "callback_url": "http://example.com/callback", "text": "Anong pangalan mo?"}, model_name="fptg__Model")
         self.model.vectorize = MagicMock(return_value=[[1, 2, 3]])
         response = self.model.respond(query)
         self.assertEqual(len(response), 1)

test/lib/model/test_image.py

@@ -25,15 +25,15 @@ def test_get_iobytes_for_image(self, mock_urlopen):
         mock_response = Mock()
         mock_response.read.return_value = image_content
         mock_urlopen.return_value = mock_response
-        image = schemas.Message(body=schemas.GenericItem(id="123", callback_url="http://example.com?callback", url="http://example.com/image.jpg"))
+        image = schemas.Message(body={"id": "123", "callback_url": "http://example.com?callback", "url": "http://example.com/image.jpg"}, model_name="audio__Model")
         result = Model().get_iobytes_for_image(image)
         self.assertIsInstance(result, io.BytesIO)
         self.assertEqual(result.read(), image_content)
 
     @patch("urllib.request.urlopen")
     def test_get_iobytes_for_image_raises_error(self, mock_urlopen):
         mock_urlopen.side_effect = URLError('test error')
-        image = schemas.Message(body=schemas.GenericItem(id="123", callback_url="http://example.com?callback", url="http://example.com/image.jpg"))
+        image = schemas.Message(body={"id": "123", "callback_url": "http://example.com?callback", "url": "http://example.com/image.jpg"}, model_name="audio__Model")
         with self.assertRaises(URLError):
             Model().get_iobytes_for_image(image)
 
@@ -42,7 +42,7 @@ def test_get_iobytes_for_image_raises_error(self, mock_urlopen):
     def test_process(self, mock_compute_pdq, mock_get_iobytes_for_image):
         mock_compute_pdq.return_value = "1001"
         mock_get_iobytes_for_image.return_value = io.BytesIO(b"image_bytes")
-        image = schemas.Message(body=schemas.GenericItem(id="123", callback_url="http://example.com?callback", url="http://example.com/image.jpg"))
+        image = schemas.Message(body={"id": "123", "callback_url": "http://example.com?callback", "url": "http://example.com/image.jpg"}, model_name="audio__Model")
         result = Model().process(image)
         self.assertEqual(result, "1001")
 

test/lib/model/test_indian_sbert.py

@@ -13,7 +13,7 @@ def setUp(self):
         self.mock_model = MagicMock()
 
     def test_vectorize(self):
-        texts = [schemas.Message(body=schemas.GenericItem(id="123", callback_url="http://example.com/callback", text="Hello, how are you?")), schemas.Message(body=schemas.GenericItem(id="123", callback_url="http://example.com/callback", text="I'm doing great, thanks!"))]
+        texts = [schemas.Message(body={"id": "123", "callback_url": "http://example.com/callback", "text": "Hello, how are you?"}, model_name="indian_sbert__Model"), schemas.Message(body={"id": "123", "callback_url": "http://example.com/callback", "text": "I'm doing great, thanks!"}, model_name="indian_sbert__Model")]
         self.model.model = self.mock_model
         self.model.model.encode = MagicMock(return_value=np.array([[4, 5, 6], [7, 8, 9]]))
         vectors = self.model.vectorize(texts)
@@ -22,7 +22,7 @@ def test_vectorize(self):
         self.assertEqual(vectors[1], [7, 8, 9])
 
     def test_respond(self):
-        query = schemas.Message(body=schemas.GenericItem(id="123", callback_url="http://example.com/callback", text="What is the capital of India?"))
+        query = schemas.Message(body={"id": "123", "callback_url": "http://example.com/callback", "text": "What is the capital of India?"}, model_name="indian_sbert__Model")
         self.model.vectorize = MagicMock(return_value=[[1, 2, 3]])
         response = self.model.respond(query)
         self.assertEqual(len(response), 1)

test/lib/model/test_meantokens.py

@@ -13,7 +13,7 @@ def setUp(self):
         self.mock_model = MagicMock()
 
     def test_vectorize(self):
-        texts = [schemas.Message(body=schemas.GenericItem(id="123", callback_url="http://example.com/callback", text="Hello, how are you?")), schemas.Message(body=schemas.GenericItem(id="123", callback_url="http://example.com/callback", text="I'm doing great, thanks!"))]
+        texts = [schemas.Message(body={"id": "123", "callback_url": "http://example.com/callback", "text": "Hello, how are you?"}, model_name="mean_tokens__Model"), schemas.Message(body={"id": "123", "callback_url": "http://example.com/callback", "text": "I'm doing great, thanks!"}, model_name="mean_tokens__Model")]
         self.model.model = self.mock_model
         self.model.model.encode = MagicMock(return_value=np.array([[4, 5, 6], [7, 8, 9]]))
         vectors = self.model.vectorize(texts)
@@ -22,7 +22,7 @@ def test_vectorize(self):
         self.assertEqual(vectors[1], [7, 8, 9])
 
     def test_respond(self):
-        query = schemas.Message(body=schemas.GenericItem(id="123", callback_url="http://example.com/callback", text="What is the capital of France?"))
+        query = schemas.Message(body={"id": "123", "callback_url": "http://example.com/callback", "text": "What is the capital of France?"}, model_name="mean_tokens__Model")
         self.model.vectorize = MagicMock(return_value=[[1, 2, 3]])
         response = self.model.respond(query)
         self.assertEqual(len(response), 1)

test/lib/model/test_video.py

@@ -32,7 +32,7 @@ def test_process_video(self, mock_pathlib, mock_upload_file_to_s3,
         mock_hash_video_output.getPureAverageFeature.return_value = "hash_value"
         mock_hash_video.return_value = mock_hash_video_output
         mock_urlopen.return_value = MagicMock(read=MagicMock(return_value=video_contents))
-        self.video_model.process(schemas.Message(body=schemas.GenericItem(id="123", callback_url="http://blah.com?callback_id=123", url="http://example.com/video.mp4")))
+        self.video_model.process(schemas.Message(body={"id": "123", "callback_url": "http://blah.com?callback_id=123", "url": "http://example.com/video.mp4"}, model_name="video__Model"))
         mock_urlopen.assert_called_once()
         mock_hash_video.assert_called_once_with(ANY, "/usr/local/bin/ffmpeg")
 
@@ -49,15 +49,15 @@ def test_tmk_program_name(self):
         self.assertEqual(result, "PrestoVideoEncoder")
 
     def test_respond_with_single_video(self):
-        video = schemas.Message(body=schemas.GenericItem(id="123", callback_url="http://blah.com?callback_id=123", url="http://example.com/video.mp4"))
+        video = schemas.Message(body={"id": "123", "callback_url": "http://blah.com?callback_id=123", "url": "http://example.com/video.mp4"}, model_name="video__Model")
         mock_process = MagicMock()
         self.video_model.process = mock_process
         result = self.video_model.respond(video)
         mock_process.assert_called_once_with(video)
         self.assertEqual(result, [video])
 
     def test_respond_with_multiple_videos(self):
-        videos = [schemas.Message(body=schemas.GenericItem(id="123", callback_url="http://blah.com?callback_id=123", url="http://example.com/video1.mp4")), schemas.Message(body=schemas.GenericItem(id="123", callback_url="http://blah.com?callback_id=123", url="http://example.com/video2.mp4"))]
+        videos = [schemas.Message(body={"id": "123", "callback_url": "http://blah.com?callback_id=123", "url": "http://example.com/video.mp4"}, model_name="video__Model"), schemas.Message(body={"id": "123", "callback_url": "http://blah.com?callback_id=123", "url": "http://example.com/video2.mp4"}, model_name="video__Model")]
         mock_process = MagicMock()
         self.video_model.process = mock_process
         result = self.video_model.respond(videos)

test/lib/queue/test_processor.py

@@ -25,7 +25,7 @@ def setUp(self, mock_get_env_setting, mock_boto_resource):
     def test_send_callbacks(self):
         # Mocking necessary methods and creating fake data
         self.queue_processor.receive_messages = MagicMock(
-            return_value=[(FakeSQSMessage(receipt_handle="blah", body=json.dumps({"body": {"callback_url": "http://example.com", "text": "This is a test", "id": 1}, "response": [1,2,3]})), self.mock_input_queue)]
+            return_value=[(FakeSQSMessage(receipt_handle="blah", body=json.dumps({"body": {"callback_url": "http://example.com", "text": "This is a test", "id": 1, "hash_value": [1,2,3]}, "model_name": "mean_tokens__Model"})), self.mock_input_queue)]
         )
         self.queue_processor.send_callback = MagicMock(return_value=None)
         self.queue_processor.delete_messages = MagicMock(return_value=None)
@@ -38,15 +38,15 @@ def test_send_callbacks(self):
 
     @patch('lib.queue.processor.requests.post')
     def test_send_callback(self, mock_post):
-        message_body = schemas.Message(body={"callback_url": "http://example.com", "text": "This is a test", "id": 123}, response=[1,2,3])
+        message_body = schemas.Message(body={"callback_url": "http://example.com", "text": "This is a test", "id": 123, "hash_value": [1,2,3]}, model_name="mean_tokens__Model")
         self.queue_processor.send_callback(message_body)
 
         mock_post.assert_called_once_with("http://example.com", json=message_body)
 
     @patch('lib.queue.processor.requests.post')
     def test_send_callback_failure(self, mock_post):
         mock_post.side_effect = Exception("Request Failed!")
-        message_body = schemas.Message(body={"callback_url": "http://example.com", "text": "This is a test", "id": 123}, response=[1,2,3])
+        message_body = schemas.Message(body={"callback_url": "http://example.com", "text": "This is a test", "id": 123, "hash_value": [1,2,3]}, model_name="mean_tokens__Model")
         with self.assertLogs(level='ERROR') as cm:
             self.queue_processor.send_callback(message_body)
         self.assertIn("Failed with Request Failed! on http://example.com with message of", cm.output[0])

test/lib/queue/test_queue.py

@@ -100,11 +100,11 @@ def test_delete_messages_from_queue(self, mock_logger):
         mock_logger.assert_called_with(f"Deleting message: {mock_messages[-1]}")
 
     def test_push_message(self):
-        message_to_push = schemas.Message(body={"id": 1, "callback_url": "http://example.com", "text": "This is a test"})
+        message_to_push = schemas.Message(body={"id": 1, "callback_url": "http://example.com", "text": "This is a test"}, model_name="mean_tokens__Model")
         # Call push_message
         returned_message = self.queue.push_message(self.queue_name_output, message_to_push)
         # Check if the message was correctly serialized and sent
-        self.mock_output_queue.send_message.assert_called_once_with(MessageBody='{"body": {"id": "1", "callback_url": "http://example.com", "url": null, "text": "This is a test", "raw": {}, "hash_value": null, "bucket": null, "outfile": null}}')
+        self.mock_output_queue.send_message.assert_called_once_with(MessageBody='{"body": {"id": "1", "callback_url": "http://example.com", "url": null, "text": "This is a test", "raw": {}, "hash_value": null}, "model_name": "mean_tokens__Model"}')
 
         self.assertEqual(returned_message, message_to_push)
 

test/lib/test_http.py

@@ -17,10 +17,10 @@ def test_process_item(self, mock_push_message, mock_create):
 
         test_data = {"id": 1, "callback_url": "http://example.com", "text": "This is a test"}
 
-        response = self.client.post("/process_item/test_process", json=test_data)
-        mock_create.assert_called_once_with("test_process")
+        response = self.client.post("/process_item/fptg__Model", json=test_data)
+        mock_create.assert_called_once_with("fptg__Model")
         self.assertEqual(response.status_code, 200)
-        self.assertEqual(response.json(), {"message": "Message pushed successfully", "queue": "test_process", "body": test_data})
+        self.assertEqual(response.json(), {"message": "Message pushed successfully", "queue": "fptg__Model", "body": test_data})
 
 
     @patch('lib.http.post_url')