117 lines
3.8 KiB
Python
117 lines
3.8 KiB
Python
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,
|
|
)
|
|
|
|
# __ __ _____ _ _ _____ _ _ _____ _
|
|
# \ \ / /\ | __ \| \ | |_ _| \ | |/ ____| |
|
|
# \ \ /\ / / \ | |__) | \| | | | | \| | | __| |
|
|
# \ \/ \/ / /\ \ | _ /| . ` | | | | . ` | | |_ | |
|
|
# \ /\ / ____ \| | \ \| |\ |_| |_| |\ | |__| |_|
|
|
# \/ \/_/ \_\_| \_\_| \_|_____|_| \_|\_____(_)
|
|
# 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
|
|
)
|