def check_content_length(self, mock_urlopen, is_header_present):
url_response = MagicMock()
if is_header_present:
url_response.info.return_value = {'Content-Length': '123'}
else:
url_response.info.return_value = {}
mock_urlopen.return_value = url_response
url_handler = server.UrlHandler('https://www.google.com')
if is_header_present:
self.assertEqual(url_handler.get_content_length(), 123)
else:
self.assertTrue(url_handler.get_content_length() is None)
def check_download_file_name(self, mock_urlopen, is_header_present, is_version_present):
url_response = MagicMock()
if is_header_present:
url_response.info.return_value = {'Content-Disposition': 'attachment; filename="test.txt"'}
else:
url_response.info.return_value = {}
mock_urlopen.return_value = url_response
url_handler = server.UrlHandler('https://www.google.com')
if is_header_present:
self.assertEqual(url_handler.get_download_file_name(), 'test.txt')
else:
if is_version_present:
self.assertEqual(url_handler.get_download_file_name('0.148'), 'presto-server-rpm-0.148.rpm')
else:
self.assertEqual(url_handler.get_download_file_name(), server.DEFAULT_RPM_NAME)