lolafect/tests/test_integration/test_connections.py

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