def test_cli_store_yaml_then_run_gcs():
"""Test storing and retrieving validation YAML when GCS env var is set."""
# Store BQ Connection
_store_bq_conn()
# Build validation and store to file
parser = cli_tools.configure_arg_parser()
mock_args = parser.parse_args(CLI_STORE_COLUMN_ARGS)
main.run(mock_args)
# Look for YAML file in GCS env directory, since that has been set
yaml_file_path = os.path.join(
os.environ[consts.ENV_DIRECTORY_VAR], "validations/", CLI_CONFIG_FILE
)
# The number of lines is not significant, except that it represents
# the exact file expected to be created. Any change to this value
# is likely to be a breaking change and must be assessed.
mgr = state_manager.StateManager()
validation_bytes = mgr._read_file(yaml_file_path)
yaml_file_str = validation_bytes.decode("utf-8")
assert len(yaml_file_str.splitlines()) == EXPECTED_NUM_YAML_LINES
# Run generated config using 'run-config' command
run_config_args = parser.parse_args(CLI_RUN_CONFIG_ARGS)
config_managers = main.build_config_managers_from_yaml(run_config_args)
main.run_validations(run_config_args, config_managers)
# Run generated config using 'configs run' command
run_config_args = parser.parse_args(CLI_CONFIGS_RUN_ARGS)
config_managers = main.build_config_managers_from_yaml(run_config_args)
main.run_validations(run_config_args, config_managers)
def test_cli_find_tables(spanner_connection_args, database_id):
_store_spanner_conn(spanner_connection_args)
parser = cli_tools.configure_arg_parser()
args = parser.parse_args(CLI_FIND_TABLES_ARGS)
tables_json = main.find_tables_using_string_matching(args)
tables = json.loads(tables_json)
assert isinstance(tables_json, str)
assert {
"schema_name": database_id,
"table_name": "array_table",
"target_schema_name": database_id,
"target_table_name": "array_table",
} in tables
assert {
"schema_name": database_id,
"table_name": "functional_alltypes",
"target_schema_name": database_id,
"target_table_name": "functional_alltypes",
} in tables
assert {
"schema_name": database_id,
"table_name": "students_pointer",
"target_schema_name": database_id,
"target_table_name": "students_pointer",
} in tables
_remove_spanner_conn()
def test_get_connection_config_from_args():
"""Test configuring arg parse in different ways."""
parser = cli_tools.configure_arg_parser()
args = parser.parse_args(CLI_ADD_CONNECTION_ARGS)
conn = cli_tools.get_connection_config_from_args(args)
assert conn["project_id"] == "example-project"
def test_configure_arg_parser_list_connections():
"""Test configuring arg parse in different ways."""
parser = cli_tools.configure_arg_parser()
args = parser.parse_args(["connections", "list"])
assert args.command == "connections"
assert args.connect_cmd == "list"
def test_cli_find_tables():
_store_bq_conn()
parser = cli_tools.configure_arg_parser()
args = parser.parse_args(CLI_FIND_TABLES_ARGS)
tables_json = main.find_tables_using_string_matching(args)
assert isinstance(tables_json, str)
assert STRING_MATCH_RESULT in tables_json
def test_configure_arg_parser_list_and_run_validation_configs():
"""Test configuring arg parse in different ways."""
parser = cli_tools.configure_arg_parser()
args = parser.parse_args(["configs", "list"])
assert args.command == "configs"
assert args.validation_config_cmd == "list"
args = parser.parse_args(["configs", "run"])
assert args.command == "configs"
assert args.validation_config_cmd == "run"
def test_create_and_list_connections(capsys, fs):
# Create Connection
parser = cli_tools.configure_arg_parser()
args = parser.parse_args(CLI_ADD_CONNECTION_ARGS)
conn = cli_tools.get_connection_config_from_args(args)
cli_tools.store_connection(args.connection_name, conn)
captured = capsys.readouterr()
assert WRITE_SUCCESS_STRING in captured.out
# List Connection
cli_tools.list_connections()
captured = capsys.readouterr()
assert captured.out == "Connection Name: test\n"
def _test_cli_yaml_local_runner(cli_args, num_yaml_lines):
"""Test storing column validation YAML."""
# Unset GCS env var so that YAML is saved locally
gcs_path = os.environ[consts.ENV_DIRECTORY_VAR]
os.environ[consts.ENV_DIRECTORY_VAR] = ""
# Store BQ Connection
_store_bq_conn()
# Build validation and store to file
parser = cli_tools.configure_arg_parser()
mock_args = parser.parse_args(cli_args)
main.run(mock_args)
yaml_file_path = CLI_CONFIG_FILE
with open(yaml_file_path, "r") as yaml_file:
assert len(yaml_file.readlines()) == num_yaml_lines
os.remove(yaml_file_path)
# Re-set GCS env var
os.environ[consts.ENV_DIRECTORY_VAR] = gcs_path
def test_cli_store_yaml_then_run_local():
"""Test storing and retrieving validation YAML when GCS env var not set."""
# Unset GCS env var so that YAML is saved locally
gcs_path = os.environ[consts.ENV_DIRECTORY_VAR]
os.environ[consts.ENV_DIRECTORY_VAR] = ""
# Store BQ Connection
_store_bq_conn()
# Build validation and store to file
parser = cli_tools.configure_arg_parser()
mock_args = parser.parse_args(CLI_STORE_COLUMN_ARGS)
main.run(mock_args)
yaml_file_path = CLI_CONFIG_FILE
with open(yaml_file_path, "r") as yaml_file:
# The number of lines is not significant, except that it represents
# the exact file expected to be created. Any change to this value
# is likely to be a breaking change and must be assessed.
assert len(yaml_file.readlines()) == EXPECTED_NUM_YAML_LINES
# Run generated config using 'run-config' command
run_config_args = parser.parse_args(CLI_RUN_CONFIG_ARGS)
config_managers = main.build_config_managers_from_yaml(run_config_args)
main.run_validations(run_config_args, config_managers)
# Run generated config using 'configs run' command
run_config_args = parser.parse_args(CLI_CONFIGS_RUN_ARGS)
config_managers = main.build_config_managers_from_yaml(run_config_args)
main.run_validations(run_config_args, config_managers)
os.remove(yaml_file_path)
# _remove_bq_conn()
# Re-set GCS env var
os.environ[consts.ENV_DIRECTORY_VAR] = gcs_path
def _store_spanner_conn(spanner_connection_args):
parser = cli_tools.configure_arg_parser()
mock_args = parser.parse_args(spanner_connection_args)
main.run_connections(mock_args)
def _store_bq_conn():
parser = cli_tools.configure_arg_parser()
mock_args = parser.parse_args(BQ_CONN_ARGS)
main.run_connections(mock_args)
def test_find_tables_config():
parser = cli_tools.configure_arg_parser()
args = parser.parse_args(CLI_FIND_TABLES_ARGS)
allowed_schemas = cli_tools.get_arg_list(args.allowed_schemas)
assert allowed_schemas[0] == "my_schema"