def test_get_response(self, mock_resp, mock_urlopen):
client = PrestoClient('any_host', 'any_user', 8080)
mock_urlopen.return_value = mock_resp
mock_resp.read.return_value = '{"message": "ok!"}'
client.get_response_from('any_uri')
self.assertEqual(client.response_from_server, {"message": "ok!"})
def test_retrieve_rows(self, mock_uri, mock_get_from_uri):
client = PrestoClient('any_host', 'any_user', 8080)
dir = os.path.abspath(os.path.dirname(__file__))
with open(dir + '/resources/valid_rest_response_level1.txt') \
as json_file:
client.response_from_server = json.load(json_file)
mock_get_from_uri.return_value = True
mock_uri.side_effect = [
"http://localhost:8080/v1/statement/2015_harih/2", ""
]
self.assertEqual(client.get_rows(), [])
self.assertEqual(client.next_uri,
"http://localhost:8080/v1/statement/2015_harih/2")
with open(dir + '/resources/valid_rest_response_level2.txt') \
as json_file:
client.response_from_server = json.load(json_file)
mock_uri.side_effect = [
"http://localhost:8080/v1/statement/2015_harih/2", ""
]
expected_row = [["uuid1", "http://localhost:8080", "presto-main:0.97",
True],
["uuid2", "http://worker:8080", "presto-main:0.97",
False]]
self.assertEqual(client.get_rows(), expected_row)
self.assertEqual(client.next_uri, "")
def test_get_response(self, mock_resp, mock_urlopen):
client = PrestoClient('any_host', 'any_user', 8080)
mock_urlopen.return_value = mock_resp
mock_resp.read.return_value = '{"message": "ok!"}'
client.get_response_from('any_uri')
self.assertEqual(client.response_from_server, {"message": "ok!"})
def test_http_call_failed(self, mock_conn):
client = PrestoClient('any_host', 'any_user', 8080)
mock_conn.side_effect = HTTPException('Error')
self.assertFalse(client.execute_query('any_sql'))
mock_conn.side_effect = socket.error('Error')
self.assertFalse(client.execute_query('any_sql'))
def test_http_call_failed(self, mock_conn):
client = PrestoClient('any_host', 'any_user', 8080)
mock_conn.side_effect = HTTPException('Error')
self.assertFalse(client.execute_query('any_sql'))
mock_conn.side_effect = socket.error('Error')
self.assertFalse(client.execute_query('any_sql'))
def test_limit_rows(self, mock_uri, mock_get_from_uri):
client = PrestoClient('any_host', 'any_user', 8080)
dir = os.path.abspath(os.path.dirname(__file__))
with open(dir + '/resources/valid_rest_response_level2.txt') \
as json_file:
client.response_from_server = json.load(json_file)
mock_get_from_uri.return_value = True
mock_uri.side_effect = ["any_next_uri", ""]
self.assertEqual(client.get_rows(0), [])
def test_limit_rows(self, mock_uri, mock_get_from_uri):
client = PrestoClient('any_host', 'any_user', 8080)
dir = os.path.abspath(os.path.dirname(__file__))
with open(dir + '/resources/valid_rest_response_level2.txt') \
as json_file:
client.response_from_server = json.load(json_file)
mock_get_from_uri.return_value = True
mock_uri.side_effect = ["any_next_uri", ""]
self.assertEqual(client.get_rows(0), [])
def test_default_request_called(self, mock_conn):
client = PrestoClient('any_host', 'any_user', 8080)
headers = {'X-Presto-Catalog': 'tpch', 'X-Presto-Schema': 'sf1',
'X-Presto-User': '******'}
client.execute_query('any_sql')
mock_conn.assert_called_with('any_host', 8080, False, URL_TIMEOUT_MS)
mock_conn().request.assert_called_with('POST', '/v1/statement',
'any_sql', headers)
self.assertTrue(mock_conn().getresponse.called)
def test_default_request_called(self, mock_conn):
client = PrestoClient('any_host', 'any_user', 8080)
headers = {
'X-Presto-Catalog': 'tpch',
'X-Presto-Schema': 'sf1',
'X-Presto-User': '******'
}
client.execute_query('any_sql')
mock_conn.assert_called_with('any_host', 8080, False, URL_TIMEOUT_MS)
mock_conn().request.assert_called_with('POST', '/v1/statement',
'any_sql', headers)
self.assertTrue(mock_conn().getresponse.called)
def test_http_answer_not_json(self, mock_response, mock_request):
client = PrestoClient('any_host', 'any_user', 8080)
mock_response.return_value.read.return_value = 'NOT JSON!'
type(mock_response.return_value).status =\
PropertyMock(return_value=200)
self.assertRaisesRegexp(ValueError, 'No JSON object could be decoded',
client.execute_query, 'any_sql')
def test_append_rows(self, mock_uri, mock_get_from_uri):
client = PrestoClient('any_host', 'any_user', 8080)
dir = os.path.abspath(os.path.dirname(__file__))
with open(dir + '/resources/valid_rest_response_level2.txt') \
as json_file:
client.response_from_server = json.load(json_file)
mock_get_from_uri.return_value = True
mock_uri.side_effect = ["any_next_uri", "any_next_next_uri", "", ""]
expected_row = [
["uuid1", "http://localhost:8080", "presto-main:0.97", True],
["uuid2", "http://worker:8080", "presto-main:0.97", False],
["uuid1", "http://localhost:8080", "presto-main:0.97", True],
["uuid2", "http://worker:8080", "presto-main:0.97", False]
]
self.assertEqual(client.get_rows(), expected_row)
def test_no_sql(self):
client = PrestoClient('any_host', 'any_user', 8080)
self.assertRaisesRegexp(
InvalidArgumentError,
'SQL query missing',
client.execute_query,
'',
)
def test_http_answer_valid(self, mock_response, mock_request):
client = PrestoClient('any_host', 'any_user', 8080)
mock_response.return_value.read.return_value = '{}'
type(mock_response.return_value).status = \
PropertyMock(return_value=200)
self.assertTrue(client.execute_query('any_sql'))
def test_http_answer_valid(self, mock_response, mock_request):
client = PrestoClient('any_host', 'any_user', 8080)
mock_response.return_value.read.return_value = '{}'
type(mock_response.return_value).status = \
PropertyMock(return_value=200)
self.assertTrue(client.execute_query('any_sql'))
def test_connection_failed(self, mock_conn):
client = PrestoClient('any_host', 'any_user', 8080)
client.execute_query('any_sql')
self.assertTrue(mock_conn().close.called)
self.assertFalse(client.execute_query('any_sql'))
def test_connection_failed(self, mock_conn):
client = PrestoClient('any_host', 'any_user', 8080)
client.execute_query('any_sql')
self.assertTrue(mock_conn().close.called)
self.assertFalse(client.execute_query('any_sql'))
def test_no_server(self):
client = PrestoClient('', 'any_user', 8080)
self.assertRaisesRegexp(InvalidArgumentError, 'Server IP missing',
client.execute_query, 'any_sql')
