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_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_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_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')