def test_error_dag_k8s(dagster_docker_image, cluster_provider):
print('--- :airflow: test_kubernetes.test_error_dag_k8s')
_check_aws_creds_available()
pipeline_name = 'demo_error_pipeline'
handle = ExecutionTargetHandle.for_pipeline_module('test_pipelines.repo',
pipeline_name)
environments_path = test_project_environments_path()
environment_yaml = [
os.path.join(environments_path, 'env_s3.yaml'),
]
environment_dict = load_yaml_from_glob_list(environment_yaml)
run_id = make_new_run_id()
execution_date = timezone.utcnow()
dag, tasks = make_airflow_dag_kubernetized_for_handle(
handle=handle,
pipeline_name=pipeline_name,
image=dagster_docker_image,
namespace='default',
environment_dict=environment_dict,
op_kwargs={
'config_file': os.environ['KUBECONFIG'],
'env_vars': {
'AWS_ACCESS_KEY_ID': os.environ['AWS_ACCESS_KEY_ID'],
'AWS_SECRET_ACCESS_KEY': os.environ['AWS_SECRET_ACCESS_KEY'],
},
},
)
with pytest.raises(AirflowException) as exc_info:
execute_tasks_in_dag(dag, tasks, run_id, execution_date)
assert 'Exception: Unusual error' in str(exc_info.value)
def test_my_custom_operator(
dagster_airflow_custom_operator_pipeline,
caplog,
): # pylint: disable=redefined-outer-name
caplog.set_level(logging.INFO, logger='CustomOperatorLogger')
pipeline_name = 'demo_pipeline'
operator = CustomOperator
environments_path = test_project_environments_path()
results = dagster_airflow_custom_operator_pipeline(
pipeline_name=pipeline_name,
handle=ReconstructableRepository.for_module('test_pipelines.repo',
pipeline_name),
operator=operator,
environment_yaml=[
os.path.join(environments_path, 'env.yaml'),
os.path.join(environments_path,
'env_filesystem_no_explicit_base_dir.yaml'),
],
)
validate_pipeline_execution(results)
log_lines = 0
for record in caplog.records:
if record.name == 'CustomOperatorLogger':
log_lines += 1
assert record.message == 'CustomOperator is called'
assert log_lines == 2
def test_s3_storage(dagster_airflow_k8s_operator_pipeline,
dagster_docker_image, cluster_provider):
print('--- :airflow: test_kubernetes.test_s3_storage')
_check_aws_creds_available()
environments_path = test_project_environments_path()
pipeline_name = 'demo_pipeline'
results = dagster_airflow_k8s_operator_pipeline(
pipeline_name=pipeline_name,
handle=ExecutionTargetHandle.for_pipeline_module(
'test_pipelines.repo', pipeline_name),
environment_yaml=[
os.path.join(environments_path, 'env.yaml'),
os.path.join(environments_path, 'env_s3.yaml'),
],
image=dagster_docker_image,
op_kwargs={
'config_file': os.environ['KUBECONFIG'],
'env_vars': {
'AWS_ACCESS_KEY_ID': os.environ['AWS_ACCESS_KEY_ID'],
'AWS_SECRET_ACCESS_KEY': os.environ['AWS_SECRET_ACCESS_KEY'],
},
},
)
validate_pipeline_execution(results)
def test_execute_on_celery( # pylint: disable=redefined-outer-name
dagster_docker_image, dagster_instance, helm_namespace
):
environment_dict = merge_dicts(
merge_yamls(
[
os.path.join(test_project_environments_path(), 'env.yaml'),
os.path.join(test_project_environments_path(), 'env_s3.yaml'),
]
),
{
'execution': {
'celery-k8s': {
'config': {
'broker': {'env': 'DAGSTER_K8S_CELERY_BROKER'},
'backend': {'env': 'DAGSTER_K8S_CELERY_BACKEND'},
'job_image': dagster_docker_image,
'job_namespace': helm_namespace,
'instance_config_map': 'dagster-instance',
'postgres_password_secret': 'dagster-postgresql-secret',
'image_pull_policy': 'Always',
'env_config_maps': ['dagster-job-runner-env'],
}
}
},
},
)
pipeline_name = 'demo_pipeline_celery'
run = dagster_instance.create_run(
pipeline_name=pipeline_name, environment_dict=environment_dict, mode='default'
)
dagster_instance.launch_run(run.run_id)
result = wait_for_job_and_get_logs(
job_name='dagster-run-%s' % run.run_id, namespace=helm_namespace
)
assert not result.get('errors')
assert result['data']
assert (
result['data']['startPipelineExecutionForCreatedRun']['__typename']
== 'StartPipelineRunSuccess'
)
def test_skip_operator(dagster_airflow_python_operator_pipeline, ): # pylint: disable=redefined-outer-name
pipeline_name = 'optional_outputs'
environments_path = test_project_environments_path()
results = dagster_airflow_python_operator_pipeline(
pipeline_name=pipeline_name,
handle=ExecutionTargetHandle.for_pipeline_module(
'test_pipelines.repo', pipeline_name),
environment_yaml=[
os.path.join(environments_path, 'env_filesystem.yaml')
],
)
validate_skip_pipeline_execution(results)
def test_gcs_storage(dagster_airflow_python_operator_pipeline, ): # pylint: disable=redefined-outer-name
pipeline_name = 'demo_pipeline_gcs'
environments_path = test_project_environments_path()
results = dagster_airflow_python_operator_pipeline(
pipeline_name=pipeline_name,
handle=ExecutionTargetHandle.for_pipeline_module(
'test_pipelines.repo', pipeline_name),
environment_yaml=[
os.path.join(environments_path, 'env.yaml'),
os.path.join(environments_path, 'env_gcs.yaml'),
],
)
validate_pipeline_execution(results)
def test_fs_storage(dagster_airflow_python_operator_pipeline, ): # pylint: disable=redefined-outer-name
pipeline_name = 'demo_pipeline'
environments_path = test_project_environments_path()
results = dagster_airflow_python_operator_pipeline(
pipeline_name=pipeline_name,
handle=ReconstructableRepository.for_module('test_pipelines.repo',
pipeline_name),
environment_yaml=[
os.path.join(environments_path, 'env.yaml'),
os.path.join(environments_path, 'env_filesystem.yaml'),
],
)
validate_pipeline_execution(results)
def test_skip_operator(dagster_airflow_docker_operator_pipeline,
dagster_docker_image): # pylint: disable=redefined-outer-name
pipeline_name = 'optional_outputs'
environments_path = test_project_environments_path()
results = dagster_airflow_docker_operator_pipeline(
pipeline_name=pipeline_name,
handle=ReconstructableRepository.for_module('test_pipelines.repo',
pipeline_name),
environment_yaml=[
os.path.join(environments_path, 'env_filesystem.yaml')
],
op_kwargs={'host_tmp_dir': '/tmp'},
image=dagster_docker_image,
)
validate_skip_pipeline_execution(results)
def test_fs_storage_no_explicit_base_dir(
dagster_airflow_docker_operator_pipeline, dagster_docker_image): # pylint: disable=redefined-outer-name
pipeline_name = 'demo_pipeline'
environments_path = test_project_environments_path()
results = dagster_airflow_docker_operator_pipeline(
pipeline_name=pipeline_name,
handle=ExecutionTargetHandle.for_pipeline_module(
'test_pipelines.repo', pipeline_name),
environment_yaml=[
os.path.join(environments_path, 'env.yaml'),
os.path.join(environments_path,
'env_filesystem_no_explicit_base_dir.yaml'),
],
image=dagster_docker_image,
)
validate_pipeline_execution(results)
def test_gcs_storage(
dagster_airflow_k8s_operator_pipeline, dagster_docker_image, cluster_provider,
):
print('--- :airflow: test_kubernetes.test_gcs_storage')
environments_path = test_project_environments_path()
pipeline_name = 'demo_pipeline_gcs'
results = dagster_airflow_k8s_operator_pipeline(
pipeline_name=pipeline_name,
handle=ReconstructableRepository.for_module('test_pipelines.repo', pipeline_name),
environment_yaml=[
os.path.join(environments_path, 'env.yaml'),
os.path.join(environments_path, 'env_gcs.yaml'),
],
image=dagster_docker_image,
op_kwargs={'config_file': os.environ['KUBECONFIG']},
)
validate_pipeline_execution(results)
def test_error_dag_containerized(dagster_docker_image): # pylint: disable=redefined-outer-name
pipeline_name = 'demo_error_pipeline'
handle = ReconstructableRepository.for_module('test_pipelines.repo',
pipeline_name)
environments_path = test_project_environments_path()
environment_yaml = [
os.path.join(environments_path, 'env_s3.yaml'),
]
environment_dict = load_yaml_from_glob_list(environment_yaml)
run_id = make_new_run_id()
execution_date = timezone.utcnow()
dag, tasks = make_airflow_dag_containerized_for_handle(
handle, pipeline_name, dagster_docker_image, environment_dict)
with pytest.raises(AirflowException) as exc_info:
execute_tasks_in_dag(dag, tasks, run_id, execution_date)
assert 'Exception: Unusual error' in str(exc_info.value)
def test_error_dag_python(): # pylint: disable=redefined-outer-name
pipeline_name = 'demo_error_pipeline'
handle = ExecutionTargetHandle.for_pipeline_module('test_pipelines.repo',
pipeline_name)
environments_path = test_project_environments_path()
environment_yaml = [
os.path.join(environments_path, 'env_filesystem.yaml'),
]
environment_dict = load_yaml_from_glob_list(environment_yaml)
execution_date = timezone.utcnow()
dag, tasks = make_airflow_dag_for_handle(handle, pipeline_name,
environment_dict)
with pytest.raises(AirflowException) as exc_info:
execute_tasks_in_dag(dag,
tasks,
run_id=make_new_run_id(),
execution_date=execution_date)
assert 'Exception: Unusual error' in str(exc_info.value)
