diff --git a/dalle3/dalle.py b/dalle3/dalle.py index 03648de..ab8d2c8 100644 --- a/dalle3/dalle.py +++ b/dalle3/dalle.py @@ -79,7 +79,7 @@ def get_time(): def get_time_save(): """Returns the current time in the format "%d-%m-%Y %H-%M-%S" """ return datetime.datetime.now().strftime("%d-%m-%Y %H-%M-%S") - + def close(self): """Closes the Chrome driver""" self.driver.quit() diff --git a/tests.py b/tests.py index 5c769cd..ffe99a7 100644 --- a/tests.py +++ b/tests.py @@ -64,38 +64,45 @@ def test_get_urls(mock_until, dalle): def mock_chrome_driver(): return Mock() + # Create a mock for the requests module @pytest.fixture def mock_requests_module(): return Mock() + # Sample cookie value for testing @pytest.fixture def sample_cookie_value(): return "sample_cookie_value" + # Sample query for testing @pytest.fixture def sample_query(): return "sample_query" + # Test the Dalle class initialization def test_dalle_init(sample_cookie_value): dalle = Dalle(sample_cookie_value) assert dalle.cookie_value == sample_cookie_value + # Test the Dalle class get_time method def test_dalle_get_time(): dalle = Dalle("") time = dalle.get_time() assert isinstance(time, str) + # Test the Dalle class get_time_save method def test_dalle_get_time_save(): dalle = Dalle("") time_save = dalle.get_time_save() assert isinstance(time_save, str) + # Test the Dalle class download method def test_dalle_download(sample_requests_module, sample_query): dalle = Dalle("") @@ -104,6 +111,7 @@ def test_dalle_download(sample_requests_module, sample_query): dalle.download(urls, save_folder) # Add assertions here to check if images were downloaded correctly + # Test the Dalle class create method def test_dalle_create(mock_chrome_driver, sample_cookie_value, sample_query): dalle = Dalle(sample_cookie_value) @@ -112,6 +120,7 @@ def test_dalle_create(mock_chrome_driver, sample_cookie_value, sample_query): assert not dalle.create(sample_query) dalle.driver.add_cookie.assert_called_once_with(cookie_value) + # Test the Dalle class get_urls method def test_dalle_get_urls(mock_chrome_driver, sample_query): dalle = Dalle("") @@ -126,6 +135,7 @@ def test_dalle_get_urls(mock_chrome_driver, sample_query): assert "https://example.com/image1.jpg" in urls assert "https://example.com/image2.jpg" in urls + # Test the Dalle class run method def test_dalle_run(mock_chrome_driver, sample_query, sample_requests_module): dalle = Dalle("") @@ -148,4 +158,78 @@ def test_dalle_close(mock_chrome_driver): mock_chrome_driver.quit.assert_called_once() +# Test the Dalle class when the create method fails +def test_dalle_create_failure(mock_chrome_driver, sample_query): + dalle = Dalle("") + dalle.driver = mock_chrome_driver + mock_chrome_driver.get.side_effect = Exception("Failed to open website") + + # Call the create method + result = dalle.create(sample_query) + + # Assertions to check if the create method returns False and if methods were called correctly + assert result is False + mock_chrome_driver.get.assert_called_with( + f"https://www.bing.com/images/create?q={sample_query}" + ) + mock_chrome_driver.add_cookie.assert_not_called() + mock_chrome_driver.refresh.assert_not_called() + + +# Test the Dalle class when the download method fails +def test_dalle_download_failure(tmpdir, sample_requests_module): + save_folder = tmpdir.join("test_images") + dalle = Dalle("") + urls = ["https://example.com/image1.jpg", "https://example.com/image2.jpg"] + sample_requests_module.get.side_effect = Exception( + "Failed to download image content" + ) + + # Call the download method + dalle.download(urls, str(save_folder)) + + # Check if the images were not saved due to download failure + assert not save_folder.join("image_1.png").check() + assert not save_folder.join("image_2.png").check() + + +# Test the Dalle class when the run method fails +def test_dalle_run_failure( + mock_chrome_driver, sample_query, sample_requests_module, tmpdir +): + save_folder = tmpdir.join("test_images") + dalle = Dalle("") + dalle.driver = mock_chrome_driver + mock_chrome_driver.get.return_value = None + mock_chrome_driver.find_elements.return_value = [ + Mock(get_attribute=Mock(return_value="https://example.com/image1.jpg")), + Mock(get_attribute=Mock(return_value="https://example.com/image2.jpg")), + ] + sample_requests_module.get.side_effect = Exception( + "Failed to download image content" + ) + + # Call the run method + download_result = dalle.run(sample_query) + + # Assertions to check if the download failed + assert download_result is False + mock_requests_module.get.assert_called_with("https://example.com/image1.jpg") + mock_requests_module.get.assert_called_with("https://example.com/image2.jpg") + + +# Test the Dalle class when the close method fails +def test_dalle_close_failure(mock_chrome_driver): + dalle = Dalle("") + dalle.driver = mock_chrome_driver + mock_chrome_driver.quit.side_effect = Exception("Failed to close Chrome driver") + + # Call the close method + dalle.close() + + # Assertions to check if the close method encountered an exception + mock_chrome_driver.quit.assert_called_once() + + +# Add more tests as needed