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