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)