import pathlib from prefect.tasks.core.function import FunctionTask from lolafect.lolaconfig import build_lolaconfig from lolafect.connections import ( connect_to_trino, close_trino_connection, _temp_secret_file_from_s3, open_ssh_tunnel_with_s3_pkey, get_local_bind_address_from_ssh_tunnel, close_ssh_tunnel, connect_to_mysql, close_mysql_connection, ) # __ __ _____ _ _ _____ _ _ _____ _ # \ \ / /\ | __ \| \ | |_ _| \ | |/ ____| | # \ \ /\ / / \ | |__) | \| | | | | \| | | __| | # \ \/ \/ / /\ \ | _ /| . ` | | | | . ` | | |_ | | # \ /\ / ____ \| | \ \| |\ |_| |_| |\ | |__| |_| # \/ \/_/ \_\_| \_\_| \_|_____|_| \_|\_____(_) # This testing suite requires: # - The calling shell to have permission in AWS # - The calling shell to be within the Mercadão network # - Do not use this tests as part of CI/CD pipelines since they are not idempotent and # rely external resources. Instead, use them manually to check yourself that things # are working properly. TEST_LOLACONFIG = build_lolaconfig(flow_name="testing-suite") def test_that_trino_connect_and_disconnect_works_properly(): connection = connect_to_trino.run( trino_credentials=TEST_LOLACONFIG.TRINO_CREDENTIALS ) connection.cursor().execute("SELECT 1") close_trino_connection.run(trino_connection=connection) def test_temporal_download_of_secret_file_works_properly_in_happy_path(): temp_file_name = "test_temp_file" with _temp_secret_file_from_s3( TEST_LOLACONFIG.S3_BUCKET_NAME, s3_file_key="env/env_prd.json", # Not a secret file, but then again, this is a test, local_temp_file_path=temp_file_name, ) as temp: temp_file_found_when_in_context_manager = pathlib.Path(temp).exists() temp_file_missing_when_outside_context_manager = not pathlib.Path( temp_file_name ).exists() assert ( temp_file_found_when_in_context_manager and temp_file_missing_when_outside_context_manager ) def test_temporal_download_of_secret_file_works_properly_even_with_exception(): temp_file_name = "test_temp_file" try: with _temp_secret_file_from_s3( TEST_LOLACONFIG.S3_BUCKET_NAME, s3_file_key="env/env_prd.json", # Not a secret file, but then again, this is a test, local_temp_file_path=temp_file_name, ) as temp: temp_file_found_when_in_context_manager = pathlib.Path(temp).exists() raise Exception # Something nasty happens within the context manager except: pass # We go with the test, ignoring the forced exception temp_file_missing_when_outside_context_manager = not pathlib.Path( temp_file_name ).exists() assert ( temp_file_found_when_in_context_manager and temp_file_missing_when_outside_context_manager ) def test_opening_and_closing_ssh_tunnel_works_properly(): test_local_bind_host = "127.0.0.1" test_local_bind_port = 12345 tunnel = open_ssh_tunnel_with_s3_pkey.run( s3_bucket_name=TEST_LOLACONFIG.S3_BUCKET_NAME, ssh_tunnel_credentials=TEST_LOLACONFIG.SSH_TUNNEL_CREDENTIALS, remote_target_host=TEST_LOLACONFIG.DW_CREDENTIALS["host"], remote_target_port=TEST_LOLACONFIG.DW_CREDENTIALS["port"], local_bind_host=test_local_bind_host, local_bind_port=test_local_bind_port, ) tunnel_was_active = tunnel.is_active local_bind_host_matches = ( get_local_bind_address_from_ssh_tunnel.run(tunnel)[0] == test_local_bind_host ) local_bind_port_matches = ( get_local_bind_address_from_ssh_tunnel.run(tunnel)[1] == test_local_bind_port ) close_ssh_tunnel.run(tunnel) tunnel_is_no_longer_active = not tunnel.is_active assert ( tunnel_was_active and tunnel_is_no_longer_active and local_bind_host_matches and local_bind_port_matches ) def test_connect_query_and_disconnect_from_mysql_with_tunnel(): test_local_bind_host = "127.0.0.1" test_local_bind_port = 12345 tunnel = open_ssh_tunnel_with_s3_pkey.run( s3_bucket_name=TEST_LOLACONFIG.S3_BUCKET_NAME, ssh_tunnel_credentials=TEST_LOLACONFIG.SSH_TUNNEL_CREDENTIALS, remote_target_host=TEST_LOLACONFIG.DW_CREDENTIALS["host"], remote_target_port=TEST_LOLACONFIG.DW_CREDENTIALS["port"], local_bind_host=test_local_bind_host, local_bind_port=test_local_bind_port, ) connection = connect_to_mysql.run( mysql_credentials=TEST_LOLACONFIG.DW_CREDENTIALS, overriding_host_and_port=get_local_bind_address_from_ssh_tunnel.run( tunnel=tunnel ), ) connection_was_open = connection.open connection.cursor().execute("SELECT 1") close_mysql_connection.run(connection=connection) close_ssh_tunnel.run(tunnel=tunnel) connection_is_closed = not connection.open assert connection_was_open and connection_is_closed