Merge pull request #1 from datarobot/andrius/address-minor-comments-f…

…rom-astronomer

Address minor comments from Astronomer.

datarobot_provider/example_dags/datarobot_pipeline_dag.py

@@ -5,8 +5,9 @@
 # This is proprietary source code of DataRobot, Inc. and its affiliates.
 #
 # Released under the terms of DataRobot Tool and Utility Agreement.
+from datetime import datetime
+
 from airflow.decorators import dag
-from airflow.utils.dates import days_ago
 
 from datarobot_provider.operators.datarobot import (
     CreateProjectOperator,
@@ -16,20 +17,13 @@
 )
 from datarobot_provider.sensors.datarobot import AutopilotCompleteSensor, ScoringCompleteSensor
 
-# These args will get passed on to each operator
-# You can override them on a per-task basis during operator initialization
-default_args = {
-    'owner': 'airflow',
-}
-
 
 @dag(
-    default_args=default_args,
     schedule_interval=None,
-    start_date=days_ago(2),
+    start_date=datetime(2022, 1, 1),
     tags=['example'],
 )
-def DataRobot_Pipeline():
+def datarobot_pipeline():
     create_project_op = CreateProjectOperator(task_id='create_project')
 
     train_models_op = TrainModelsOperator(
@@ -63,4 +57,4 @@ def DataRobot_Pipeline():
     )
 
 
-datarobot_pipeline = DataRobot_Pipeline()
+datarobot_pipeline_dag = datarobot_pipeline()

tests/unit/dags/test_datarobot_pipeline_dag.py

@@ -8,7 +8,7 @@
 import pytest
 from airflow.models import DagBag
 
-from datarobot_provider.example_dags.datarobot_pipeline_dag import DataRobot_Pipeline
+from datarobot_provider.example_dags.datarobot_pipeline_dag import datarobot_pipeline
 
 
 @pytest.fixture()
@@ -17,7 +17,7 @@ def dagbag(provider_dir):
 
 
 def test_dag_loaded(dagbag):
-    dag = dagbag.get_dag(dag_id="DataRobot_Pipeline")
+    dag = dagbag.get_dag(dag_id="datarobot_pipeline")
     assert dagbag.import_errors == {}
     assert dag is not None
     assert len(dag.tasks) == 6
@@ -32,7 +32,7 @@ def assert_dag_dict_equal(source, dag):
 
 
 def test_dag_structure():
-    dag = DataRobot_Pipeline()
+    dag = datarobot_pipeline()
     assert_dag_dict_equal(
         {
             "create_project": ["train_models", "check_autopilot_complete", "deploy_recommended_model"],