Spaces:
No application file
No application file
| import unittest | |
| import os | |
| from PIL import Image | |
| from unittest.mock import patch, Mock | |
| from tests import check_url, load_image, get_caption | |
| class TestWebLinksHandler(unittest.TestCase): | |
| def test_check_url_valid(self): | |
| valid_url = "https://www.example.com" | |
| self.assertTrue(check_url(valid_url)) | |
| def test_check_url_invalid(self): | |
| invalid_url = "not_a_valid_url" | |
| self.assertFalse(check_url(invalid_url)) | |
| def test_load_image_from_url(self, mock_get): | |
| url = "https://www.example.com/image.jpg" | |
| response_mock = Mock() | |
| response_mock.raw = Mock() | |
| mock_get.return_value = response_mock | |
| image = load_image(url) | |
| self.assertIsInstance(image, Image.Image) | |
| def test_load_image_from_file(self): | |
| file_path = "path/to/local/image.jpg" | |
| # You may want to create a temporary image for testing purposes | |
| # and delete it after the test is run. | |
| open(file_path, 'w').close() | |
| image = load_image(file_path) | |
| self.assertIsInstance(image, Image.Image) | |
| def test_get_caption(self, mock_open): | |
| # Mocking Image processor, model, and tokenizer | |
| mock_image_processor = Mock() | |
| mock_model = Mock() | |
| mock_tokenizer = Mock() | |
| # Mocking Image object | |
| mock_image = Mock() | |
| mock_image_processor.return_value = {'image_key': mock_image} | |
| # Mocking the model.generate method | |
| mock_model.generate.return_value = torch.tensor([[1, 2, 3]]) | |
| # Mocking tokenizer.batch_decode method | |
| mock_tokenizer.batch_decode.return_value = ['test caption'] | |
| with patch('builtins.print') as mock_print: | |
| caption = get_caption(mock_model, mock_image_processor, mock_tokenizer, 'dummy_image_path') | |
| mock_open.assert_called_once() | |
| mock_model.generate.assert_called_once_with(**{'image_key': mock_image}) | |
| mock_tokenizer.batch_decode.assert_called_once_with(torch.tensor([[1, 2, 3]]), skip_special_tokens=True) | |
| mock_print.assert_called_once_with('test caption') | |
| self.assertEqual(caption, 'test caption') | |
| if __name__ == '__main__': | |
| unittest.main() | |