def test_connection_fail(self):
test_executor = OpenSearchConnection(endpoint=INVALID_ENDPOINT)
err_message = "Can not connect to endpoint %s" % INVALID_ENDPOINT
with mock.patch("sys.exit") as mock_sys_exit, mock.patch("src.opensearch_sql_cli.opensearch_connection.click.secho") as mock_secho:
test_executor.set_connection()
mock_sys_exit.assert_called()
mock_secho.assert_called_with(message=err_message, fg="red")
def test_get_od_client(self):
od_test_executor = OpenSearchConnection(endpoint=OPEN_DISTRO_ENDPOINT, http_auth=AUTH)
with mock.patch.object(OpenSearch, "__init__", return_value=None) as mock_es:
od_test_executor.get_open_distro_client()
mock_es.assert_called_with(
[OPEN_DISTRO_ENDPOINT], http_auth=AUTH, verify_certs=False, ssl_context=od_test_executor.ssl_context,
connection_class=RequestsHttpConnection
)
def test_get_aes_client(self):
aes_test_executor = OpenSearchConnection(endpoint=AES_ENDPOINT, use_aws_authentication=True)
with mock.patch.object(OpenSearch, "__init__", return_value=None) as mock_es:
aes_test_executor.get_aes_client()
mock_es.assert_called_with(
hosts=[AES_ENDPOINT],
http_auth=aes_test_executor.aws_auth,
use_ssl=True,
verify_certs=True,
connection_class=RequestsHttpConnection,
)
def test_lost_connection(self):
test_esexecutor = OpenSearchConnection(endpoint=INVALID_ENDPOINT)
def side_effect_set_connection(is_reconnected):
if is_reconnected:
pass
else:
return ConnectionError()
with mock.patch(
"src.opensearch_sql_cli.opensearch_connection.click.secho"
) as mock_secho, mock.patch.object(
test_esexecutor, "set_connection") as mock_set_connection:
# Assume reconnection success
mock_set_connection.side_effect = side_effect_set_connection(
is_reconnected=True)
test_esexecutor.handle_server_close_connection()
mock_secho.assert_any_call(message="Reconnecting...", fg="green")
mock_secho.assert_any_call(
message="Reconnected! Please run query again", fg="green")
# Assume reconnection fail
mock_set_connection.side_effect = side_effect_set_connection(
is_reconnected=False)
test_esexecutor.handle_server_close_connection()
mock_secho.assert_any_call(message="Reconnecting...", fg="green")
mock_secho.assert_any_call(
message=
"Connection Failed. Check your OpenSearch is running and then come back",
fg="red")
def test_select_client(self):
od_test_executor = OpenSearchConnection(endpoint=OPEN_DISTRO_ENDPOINT, http_auth=AUTH)
aes_test_executor = OpenSearchConnection(endpoint=AES_ENDPOINT, use_aws_authentication=True)
with mock.patch.object(od_test_executor, "get_open_distro_client") as mock_od_client, mock.patch.object(
OpenSearchConnection, "is_sql_plugin_installed", return_value=True
):
od_test_executor.set_connection()
mock_od_client.assert_called()
with mock.patch.object(aes_test_executor, "get_aes_client") as mock_aes_client, mock.patch.object(
OpenSearchConnection, "is_sql_plugin_installed", return_value=True
):
aes_test_executor.set_connection()
mock_aes_client.assert_called()
def test_reconnection_exception(self):
test_executor = OpenSearchConnection(endpoint=INVALID_ENDPOINT)
with pytest.raises(ConnectionError) as error:
assert test_executor.set_connection(True)
def get_connection():
test_es_connection = OpenSearchConnection(endpoint=ENDPOINT)
test_es_connection.set_connection(is_reconnect=True)
return test_es_connection