diff --git a/test/test_image_to_video_service.py b/test/test_image_to_video_service.py new file mode 100644 index 0000000..87e08f6 --- /dev/null +++ b/test/test_image_to_video_service.py @@ -0,0 +1,127 @@ +import unittest +from unittest.mock import MagicMock, patch +import json +import sys +import os +from PIL import Image +import io + +# Add root to the path so we can import src as a package +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) + +from src.image_to_video_service import ImageToVideoService, VideoResult +from src.config import Config + +class TestImageToVideoService(unittest.TestCase): + + def setUp(self): + self.api_key = "test_api_key" + self.service = ImageToVideoService(api_key=self.api_key) + + # Patch Config + self.config_patcher = patch('src.image_to_video_service.Config') + self.mock_config = self.config_patcher.start() + self.mock_config.DEFAULT_STYLE = "" + self.mock_config.DEFAULT_DURATION = 5 + self.mock_config.DEFAULT_AUDIO_ENABLED = True + self.mock_config.DURATION_OPTIONS = [5, 10] + self.mock_config.REQUEST_TIMEOUT = 30 + self.mock_config.KEYFRAME_POLLING_INTERVAL = 0.1 + self.mock_config.KEYFRAME_MAX_POLL_TIME = 1 + self.mock_config.get_image_to_video_models.return_value = ["wan2.5-i2v-preview"] + self.mock_config.get_max_prompt_length.return_value = 1000 + self.mock_config.get_supported_resolutions_for_model.return_value = ["1080P"] + self.mock_config.supports_audio.return_value = True + self.mock_config.validate_image_upload.return_value = None # Valid image + + # Mock oss_service + self.oss_patcher = patch('src.image_to_video_service.oss_service') + self.mock_oss = self.oss_patcher.start() + + def tearDown(self): + self.config_patcher.stop() + self.oss_patcher.stop() + + def create_dummy_image(self): + img = Image.new('RGB', (100, 100), color = 'red') + img_byte_arr = io.BytesIO() + img.save(img_byte_arr, format='JPEG') + img_byte_arr.seek(0) + # Mock file-like object attributes + img_byte_arr.size = 1024 # Fake size + return img_byte_arr + + @patch('requests.post') + @patch('src.image_to_video_service.ImageToVideoService._poll_task_result') + @patch('src.image_to_video_service.ImageToVideoService._download_video_locally') + def test_generate_video_success(self, mock_download, mock_poll, mock_post): + # Setup mocks + self.mock_oss.upload_image.return_value = ("http://oss.example.com/image.jpg", {"size": 1024}) + + mock_post.return_value.status_code = 200 + mock_post.return_value.json.return_value = { + 'output': {'task_id': 'task_123'} + } + + mock_poll.return_value = "http://example.com/video.mp4" + mock_download.return_value = "/tmp/video.mp4" + + # Execute + image_file = self.create_dummy_image() + result = self.service.generate_video( + image_file=image_file, + prompt="A moving cat", + model="wan2.5-i2v-preview" + ) + + # Verify + self.assertTrue(result.success) + self.assertEqual(result.video_url, "http://example.com/video.mp4") + self.assertEqual(result.task_id, "task_123") + + # Verify OSS upload + self.mock_oss.upload_image.assert_called_once() + + # Verify API call + mock_post.assert_called_once() + args, kwargs = mock_post.call_args + self.assertEqual(kwargs['json']['input']['img_url'], "http://oss.example.com/image.jpg") + self.assertEqual(kwargs['json']['input']['prompt'], "A moving cat") + + def test_validate_inputs_no_image(self): + result = self.service.generate_video(image_file=None) + self.assertFalse(result.success) + self.assertEqual(result.error_message, "Image file is required") + + def test_validate_inputs_invalid_model(self): + image_file = self.create_dummy_image() + result = self.service.generate_video(image_file=image_file, model="invalid-model") + self.assertFalse(result.success) + self.assertTrue("Invalid model" in result.error_message) + + @patch('requests.post') + def test_upload_failure(self, mock_post): + self.mock_oss.upload_image.return_value = (None, None) + + image_file = self.create_dummy_image() + result = self.service.generate_video(image_file=image_file) + + self.assertFalse(result.success) + self.assertTrue("Failed to process and upload image" in result.error_message) + mock_post.assert_not_called() + + @patch('requests.post') + def test_submit_task_failure(self, mock_post): + self.mock_oss.upload_image.return_value = ("http://oss.example.com/image.jpg", {}) + + mock_post.return_value.status_code = 400 + mock_post.return_value.json.return_value = {"message": "Bad Request"} + + image_file = self.create_dummy_image() + result = self.service.generate_video(image_file=image_file) + + self.assertFalse(result.success) + self.assertTrue("Bad Request" in result.error_message) + +if __name__ == '__main__': + unittest.main() diff --git a/test/test_keyframe_to_video_service.py b/test/test_keyframe_to_video_service.py new file mode 100644 index 0000000..c955653 --- /dev/null +++ b/test/test_keyframe_to_video_service.py @@ -0,0 +1,130 @@ +import unittest +from unittest.mock import MagicMock, patch +import json +import sys +import os +from PIL import Image +import io + +# Add root to the path so we can import src as a package +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) + +from src.keyframe_to_video_service import KeyFrameVideoService, VideoResult +from src.config import Config + +class TestKeyFrameVideoService(unittest.TestCase): + + def setUp(self): + self.api_key = "test_api_key" + self.service = KeyFrameVideoService(api_key=self.api_key) + + # Patch Config + self.config_patcher = patch('src.keyframe_to_video_service.Config') + self.mock_config = self.config_patcher.start() + self.mock_config.DEFAULT_STYLE = "" + self.mock_config.REQUEST_TIMEOUT = 30 + self.mock_config.KEYFRAME_POLLING_INTERVAL = 0.1 + self.mock_config.KEYFRAME_MAX_POLL_TIME = 1 + self.mock_config.MAX_PROMPT_LENGTH = 1000 + self.mock_config.get_keyframe_to_video_models.return_value = ["wanx2.1-kf2v-plus"] + self.mock_config.validate_image_upload.return_value = None # Valid image + + # Mock oss_service + self.oss_patcher = patch('src.keyframe_to_video_service.oss_service') + self.mock_oss = self.oss_patcher.start() + + def tearDown(self): + self.config_patcher.stop() + self.oss_patcher.stop() + + def create_dummy_image(self): + img = Image.new('RGB', (100, 100), color = 'red') + img_byte_arr = io.BytesIO() + img.save(img_byte_arr, format='JPEG') + img_byte_arr.seek(0) + img_byte_arr.size = 1024 + return img_byte_arr + + @patch('requests.post') + @patch('src.keyframe_to_video_service.KeyFrameVideoService._poll_task_result') + @patch('src.keyframe_to_video_service.KeyFrameVideoService._download_video_locally') + def test_generate_video_success(self, mock_download, mock_poll, mock_post): + # Setup mocks + self.mock_oss.upload_image.side_effect = [ + ("http://oss.example.com/start.jpg", {"size": 1024}), + ("http://oss.example.com/end.jpg", {"size": 1024}) + ] + + mock_post.return_value.status_code = 200 + mock_post.return_value.json.return_value = { + 'output': {'task_id': 'task_123'} + } + + mock_poll.return_value = "http://example.com/video.mp4" + mock_download.return_value = "/tmp/video.mp4" + + # Execute + start_img = self.create_dummy_image() + end_img = self.create_dummy_image() + + result = self.service.generate_video( + start_frame_file=start_img, + end_frame_file=end_img, + prompt="Transition", + model="wanx2.1-kf2v-plus" + ) + + # Verify + self.assertTrue(result.success) + self.assertEqual(result.video_url, "http://example.com/video.mp4") + self.assertEqual(result.task_id, "task_123") + + # Verify OSS upload (called twice) + self.assertEqual(self.mock_oss.upload_image.call_count, 2) + + # Verify API call + mock_post.assert_called_once() + args, kwargs = mock_post.call_args + self.assertEqual(kwargs['json']['input']['first_frame_url'], "http://oss.example.com/start.jpg") + self.assertEqual(kwargs['json']['input']['last_frame_url'], "http://oss.example.com/end.jpg") + self.assertEqual(kwargs['json']['input']['prompt'], "Transition") + + def test_validate_inputs_missing_frames(self): + img = self.create_dummy_image() + + # Missing start frame + result = self.service.generate_video(start_frame_file=None, end_frame_file=img) + self.assertFalse(result.success) + self.assertTrue("Start frame image is required" in result.error_message) + + # Missing end frame + result = self.service.generate_video(start_frame_file=img, end_frame_file=None) + self.assertFalse(result.success) + self.assertTrue("End frame image is required" in result.error_message) + + def test_validate_inputs_invalid_model(self): + img = self.create_dummy_image() + result = self.service.generate_video( + start_frame_file=img, + end_frame_file=img, + model="invalid-model" + ) + self.assertFalse(result.success) + self.assertTrue("Invalid model" in result.error_message) + + @patch('requests.post') + def test_upload_failure(self, mock_post): + # First upload fails + self.mock_oss.upload_image.side_effect = [(None, None), (None, None)] + + start_img = self.create_dummy_image() + end_img = self.create_dummy_image() + + result = self.service.generate_video(start_frame_file=start_img, end_frame_file=end_img) + + self.assertFalse(result.success) + self.assertTrue("Failed to process and upload" in result.error_message) + mock_post.assert_not_called() + +if __name__ == '__main__': + unittest.main() diff --git a/test/test_text_to_video_service.py b/test/test_text_to_video_service.py new file mode 100644 index 0000000..2e9cad5 --- /dev/null +++ b/test/test_text_to_video_service.py @@ -0,0 +1,120 @@ +import unittest +from unittest.mock import MagicMock, patch +import json +import sys +import os + +# Add root to the path so we can import src as a package +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) + +from src.text_to_video_service import TextToVideoService, VideoResult +from src.config import Config + +class TestTextToVideoService(unittest.TestCase): + + def setUp(self): + self.api_key = "test_api_key" + self.service = TextToVideoService(api_key=self.api_key) + + # Patch Config to have predictable values + self.config_patcher = patch('src.text_to_video_service.Config') + self.mock_config = self.config_patcher.start() + self.mock_config.DEFAULT_STYLE = "" + self.mock_config.DEFAULT_ASPECT_RATIO = "16:9" + self.mock_config.DEFAULT_MODEL = "wan2.5-t2v-preview" + self.mock_config.DEFAULT_DURATION = 5 + self.mock_config.DEFAULT_AUDIO_ENABLED = True + self.mock_config.STYLE_OPTIONS = ["", "Cinematic"] + self.mock_config.ASPECT_RATIO_OPTIONS = ["16:9", "1:1"] + self.mock_config.DURATION_OPTIONS = [5, 10] + self.mock_config.REQUEST_TIMEOUT = 30 + self.mock_config.POLLING_INTERVAL = 0.1 # Fast polling for tests + self.mock_config.MAX_POLL_TIME = 1 + + # Setup model options for validation + self.mock_config.get_text_to_video_models.return_value = ["wan2.5-t2v-preview"] + self.mock_config.get_max_prompt_length.return_value = 1000 + self.mock_config.get_supported_resolutions_for_model.return_value = ["1080P"] + self.mock_config.uses_resolution_label.return_value = True + self.mock_config.supports_audio.return_value = True + + def tearDown(self): + self.config_patcher.stop() + + @patch('requests.post') + @patch('src.text_to_video_service.TextToVideoService._poll_task_result') + @patch('src.text_to_video_service.TextToVideoService._download_video_locally') + def test_generate_video_success(self, mock_download, mock_poll, mock_post): + # Setup mocks + mock_post.return_value.status_code = 200 + mock_post.return_value.json.return_value = { + 'output': {'task_id': 'task_123'} + } + + mock_poll.return_value = "http://example.com/video.mp4" + mock_download.return_value = "/tmp/video.mp4" + + # Execute + result = self.service.generate_video( + prompt="A cute cat", + style="Cinematic", + aspect_ratio="16:9", + model="wan2.5-t2v-preview" + ) + + # Verify + self.assertTrue(result.success) + self.assertEqual(result.video_url, "http://example.com/video.mp4") + self.assertEqual(result.local_video_path, "/tmp/video.mp4") + self.assertEqual(result.task_id, "task_123") + + # Verify API call + mock_post.assert_called_once() + args, kwargs = mock_post.call_args + self.assertEqual(kwargs['json']['input']['prompt'], "A cute cat") + self.assertEqual(kwargs['json']['parameters']['resolution'], "1080P") + + def test_validate_inputs_empty_prompt(self): + result = self.service.generate_video(prompt="") + self.assertFalse(result.success) + self.assertEqual(result.error_message, "Prompt cannot be empty") + + def test_validate_inputs_invalid_style(self): + result = self.service.generate_video(prompt="test", style="InvalidStyle") + self.assertFalse(result.success) + self.assertTrue("Invalid style" in result.error_message) + + def test_validate_inputs_invalid_model(self): + self.mock_config.get_text_to_video_models.return_value = ["wan2.5-t2v-preview"] + result = self.service.generate_video(prompt="test", model="invalid-model") + self.assertFalse(result.success) + self.assertTrue("Invalid model" in result.error_message) + + @patch('requests.post') + def test_submit_task_failure(self, mock_post): + mock_post.return_value.status_code = 400 + mock_post.return_value.json.return_value = {"message": "Bad Request"} + + result = self.service.generate_video(prompt="test") + + self.assertFalse(result.success) + self.assertTrue("Bad Request" in result.error_message) + + @patch('requests.post') + @patch('src.text_to_video_service.TextToVideoService._poll_task_result') + def test_generate_video_poll_failure(self, mock_poll, mock_post): + mock_post.return_value.status_code = 200 + mock_post.return_value.json.return_value = { + 'output': {'task_id': 'task_123'} + } + + mock_poll.return_value = None # Polling failed or timed out + + result = self.service.generate_video(prompt="test") + + self.assertFalse(result.success) + self.assertEqual(result.task_id, "task_123") + self.assertTrue("Video generation failed or timed out" in result.error_message) + +if __name__ == '__main__': + unittest.main()