def test_insert_document(self, cosmos_mock):
test_id = str(uuid.uuid4())
cosmos_mock.return_value.CreateItem.return_value = {'id': test_id}
self.cosmos = AzureCosmosInsertDocumentOperator(
database_name=self.test_database_name,
collection_name=self.test_collection_name,
document={
'id': test_id,
'data': 'sometestdata'
},
azure_cosmos_conn_id='azure_cosmos_test_key_id',
task_id='azure_cosmos_sensor')
expected_calls = [
mock.call().CreateItem(
'dbs/' + self.test_database_name + '/colls/' +
self.test_collection_name, {
'data': 'sometestdata',
'id': test_id
})
]
self.cosmos.execute(None)
cosmos_mock.assert_any_call(self.test_end_point,
{'masterKey': self.test_master_key})
cosmos_mock.assert_has_calls(expected_calls)
class TestAzureCosmosDbHook(unittest.TestCase):
# Set up an environment to test with
def setUp(self):
# set up some test variables
self.test_end_point = 'https://*****:*****@mock.patch('azure.cosmos.cosmos_client.CosmosClient')
def test_insert_document(self, cosmos_mock):
test_id = str(uuid.uuid4())
cosmos_mock.return_value.CreateItem.return_value = {'id': test_id}
self.cosmos = AzureCosmosInsertDocumentOperator(
database_name=self.test_database_name,
collection_name=self.test_collection_name,
document={
'id': test_id,
'data': 'sometestdata'
},
azure_cosmos_conn_id='azure_cosmos_test_key_id',
task_id='azure_cosmos_sensor')
expected_calls = [
mock.call().CreateItem(
'dbs/' + self.test_database_name + '/colls/' +
self.test_collection_name, {
'data': 'sometestdata',
'id': test_id
})
]
self.cosmos.execute(None)
cosmos_mock.assert_any_call(self.test_end_point,
{'masterKey': self.test_master_key})
cosmos_mock.assert_has_calls(expected_calls)
class TestAzureCosmosDbHook(unittest.TestCase):
# Set up an environment to test with
def setUp(self):
# set up some test variables
self.test_end_point = 'https://*****:*****@mock.patch('azure.cosmos.cosmos_client.CosmosClient')
def test_insert_document(self, cosmos_mock):
test_id = str(uuid.uuid4())
cosmos_mock.return_value.CreateItem.return_value = {'id': test_id}
self.cosmos = AzureCosmosInsertDocumentOperator(
database_name=self.test_database_name,
collection_name=self.test_collection_name,
document={'id': test_id, 'data': 'sometestdata'},
azure_cosmos_conn_id='azure_cosmos_test_key_id',
task_id='azure_cosmos_sensor')
expected_calls = [mock.call().CreateItem(
'dbs/' + self.test_database_name + '/colls/' + self.test_collection_name,
{'data': 'sometestdata', 'id': test_id})]
self.cosmos.execute(None)
cosmos_mock.assert_any_call(self.test_end_point, {'masterKey': self.test_master_key})
cosmos_mock.assert_has_calls(expected_calls)
def test_insert_document(self, cosmos_mock):
test_id = str(uuid.uuid4())
cosmos_mock.return_value.CreateItem.return_value = {'id': test_id}
self.cosmos = AzureCosmosInsertDocumentOperator(
database_name=self.test_database_name,
collection_name=self.test_collection_name,
document={'id': test_id, 'data': 'sometestdata'},
azure_cosmos_conn_id='azure_cosmos_test_key_id',
task_id='azure_cosmos_sensor')
expected_calls = [mock.call().CreateItem(
'dbs/' + self.test_database_name + '/colls/' + self.test_collection_name,
{'data': 'sometestdata', 'id': test_id})]
self.cosmos.execute(None)
cosmos_mock.assert_any_call(self.test_end_point, {'masterKey': self.test_master_key})
cosmos_mock.assert_has_calls(expected_calls)
'start_date': dates.days_ago(2),
'email': ['*****@*****.**'],
'email_on_failure': False,
'email_on_retry': False
}
dag = DAG('example_azure_cosmosdb_sensor', default_args=default_args)
dag.doc_md = __doc__
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',
dag=dag)
t2 = AzureCosmosInsertDocumentOperator(
task_id='insert_cosmos_file',
dag=dag,
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