Python AzureCosmosDocumentSensor Examples

Example #1

 def test_should_return_false_on_no_document(self, mock_hook):
     mock_instance = mock_hook.return_value
     mock_instance.get_document.return_value = None  # Indicate document not returned
     sensor = AzureCosmosDocumentSensor(
         task_id="test-task-2",
         database_name=DB_NAME,
         collection_name=COLLECTION_NAME,
         document_id=DOCUMENT_ID,
     )
     result = sensor.poke(None)
     mock_instance.get_document.assert_called_once_with(DOCUMENT_ID, DB_NAME, COLLECTION_NAME)
     self.assertEqual(result, False)

Example #2

File: test_azure_cosmos.py Project: ysktir/airflow-1

 def test_should_call_hook_with_args(self, mock_hook):
     mock_instance = mock_hook.return_value
     mock_instance.get_document.return_value = True  # Indicate document returned
     sensor = AzureCosmosDocumentSensor(
         task_id="test-task-1",
         database_name=DB_NAME,
         collection_name=COLLECTION_NAME,
         document_id=DOCUMENT_ID,
     )
     result = sensor.poke(None)
     mock_instance.get_document.assert_called_once_with(
         DOCUMENT_ID, DB_NAME, COLLECTION_NAME)
     assert result is True

Example #3

File: example_azure_cosmosdb.py Project: bravefoot/incubator-airflow

    dag_id='example_azure_cosmosdb_sensor',
    default_args={
        'owner': 'airflow',
        'depends_on_past': False,
        'email': ['*****@*****.**'],
        'email_on_failure': False,
        'email_on_retry': False,
    },
    start_date=dates.days_ago(2),
    doc_md=__doc__,
    tags=['example'],
) as dag:

    t1 = AzureCosmosDocumentSensor(
        task_id='check_cosmos_file',
        database_name='airflow_example_db',
        collection_name='airflow_example_coll',
        document_id='airflow_checkid',
        azure_cosmos_conn_id='azure_cosmos_default',
    )

    t2 = AzureCosmosInsertDocumentOperator(
        task_id='insert_cosmos_file',
        database_name='airflow_example_db',
        collection_name='new-collection',
        document={"id": "someuniqueid", "param1": "value1", "param2": "value2"},
        azure_cosmos_conn_id='azure_cosmos_default',
    )

    t1 >> t2

Example #4

File: example_azure_cosmosdb.py Project: yulei-li/airflow

"""

from datetime import datetime

from airflow import DAG
from airflow.providers.microsoft.azure.operators.azure_cosmos import AzureCosmosInsertDocumentOperator
from airflow.providers.microsoft.azure.sensors.azure_cosmos import AzureCosmosDocumentSensor

with DAG(
    dag_id='example_azure_cosmosdb_sensor',
    default_args={'database_name': 'airflow_example_db'},
    start_date=datetime(2021, 1, 1),
    catchup=False,
    doc_md=__doc__,
    tags=['example'],
) as dag:

    t1 = AzureCosmosDocumentSensor(
        task_id='check_cosmos_file',
        collection_name='airflow_example_coll',
        document_id='airflow_checkid',
    )

    t2 = AzureCosmosInsertDocumentOperator(
        task_id='insert_cosmos_file',
        collection_name='new-collection',
        document={"id": "someuniqueid", "param1": "value1", "param2": "value2"},
    )

    t1 >> t2