def test_upsert_document(self, cosmos_mock):
test_id = str(uuid.uuid4())
cosmos_mock.return_value.CreateItem.return_value = {'id': test_id}
self.cosmos = AzureCosmosDBHook(
azure_cosmos_conn_id='azure_cosmos_test_key_id')
returned_item = self.cosmos.upsert_document(
{'data1': 'somedata'},
database_name=self.test_database_name,
collection_name=self.test_collection_name,
document_id=test_id)
expected_calls = [
mock.call().CreateItem(
'dbs/' + self.test_database_name + '/colls/' +
self.test_collection_name, {
'data1': 'somedata',
'id': test_id
})
]
cosmos_mock.assert_any_call(self.test_end_point,
{'masterKey': self.test_master_key})
cosmos_mock.assert_has_calls(expected_calls)
logging.getLogger().info(returned_item)
self.assertEqual(returned_item['id'], test_id)
def test_create_container_default(self, cosmos_mock):
self.cosmos = AzureCosmosDBHook(azure_cosmos_conn_id='azure_cosmos_test_key_id')
self.cosmos.create_collection(self.test_collection_name)
expected_calls = [mock.call().CreateContainer(
'dbs/test_database_default',
{'id': self.test_collection_name})]
cosmos_mock.assert_any_call(self.test_end_point, {'masterKey': self.test_master_key})
cosmos_mock.assert_has_calls(expected_calls)
def test_delete_database(self, cosmos_mock):
self.cosmos = AzureCosmosDBHook(
azure_cosmos_conn_id='azure_cosmos_test_key_id')
self.cosmos.delete_database(self.test_database_name)
expected_calls = [mock.call().DeleteDatabase('dbs/test_database_name')]
cosmos_mock.assert_any_call(self.test_end_point,
{'masterKey': self.test_master_key})
cosmos_mock.assert_has_calls(expected_calls)
def test_create_database(self, mock_cosmos):
hook = AzureCosmosDBHook(
azure_cosmos_conn_id='azure_cosmos_test_key_id')
hook.create_database(self.test_database_name)
expected_calls = [
mock.call().CreateDatabase({'id': self.test_database_name})
]
mock_cosmos.assert_any_call(self.test_end_point,
{'masterKey': self.test_master_key})
mock_cosmos.assert_has_calls(expected_calls)
def test_delete_container_default(self, mock_cosmos):
hook = AzureCosmosDBHook(
azure_cosmos_conn_id='azure_cosmos_test_key_id')
hook.delete_collection(self.test_collection_name)
expected_calls = [
mock.call().DeleteContainer(
'dbs/test_database_default/colls/test_collection_name')
]
mock_cosmos.assert_any_call(self.test_end_point,
{'masterKey': self.test_master_key})
mock_cosmos.assert_has_calls(expected_calls)
def execute(self, context):
# Create the hook
hook = AzureCosmosDBHook(azure_cosmos_conn_id=self.azure_cosmos_conn_id)
# Create the DB if it doesn't already exist
if not hook.does_database_exist(self.database_name):
hook.create_database(self.database_name)
# Create the collection as well
if not hook.does_collection_exist(self.collection_name, self.database_name):
hook.create_collection(self.collection_name, self.database_name)
# finally insert the document
hook.upsert_document(self.document, self.database_name, self.collection_name)
def test_upsert_document_default(self, mock_cosmos):
test_id = str(uuid.uuid4())
mock_cosmos.return_value.CreateItem.return_value = {'id': test_id}
hook = AzureCosmosDBHook(
azure_cosmos_conn_id='azure_cosmos_test_key_id')
returned_item = hook.upsert_document({'id': test_id})
expected_calls = [
mock.call().CreateItem(
'dbs/' + self.test_database_default + '/colls/' +
self.test_collection_default, {'id': test_id})
]
mock_cosmos.assert_any_call(self.test_end_point,
{'masterKey': self.test_master_key})
mock_cosmos.assert_has_calls(expected_calls)
logging.getLogger().info(returned_item)
self.assertEqual(returned_item['id'], test_id)
def test_create_container_default(self, cosmos_mock):
self.cosmos = AzureCosmosDBHook(azure_cosmos_conn_id='azure_cosmos_test_key_id')
self.cosmos.create_collection(self.test_collection_name)
expected_calls = [mock.call().CreateContainer(
'dbs/test_database_default',
{'id': self.test_collection_name})]
cosmos_mock.assert_any_call(self.test_end_point, {'masterKey': self.test_master_key})
cosmos_mock.assert_has_calls(expected_calls)
def test_insert_documents(self, cosmos_mock):
test_id1 = str(uuid.uuid4())
test_id2 = str(uuid.uuid4())
test_id3 = str(uuid.uuid4())
documents = [{
'id': test_id1,
'data': 'data1'
}, {
'id': test_id2,
'data': 'data2'
}, {
'id': test_id3,
'data': 'data3'
}]
self.cosmos = AzureCosmosDBHook(
azure_cosmos_conn_id='azure_cosmos_test_key_id')
returned_item = self.cosmos.insert_documents(documents)
expected_calls = [
mock.call().CreateItem(
'dbs/' + self.test_database_default + '/colls/' +
self.test_collection_default, {
'data': 'data1',
'id': test_id1
}),
mock.call().CreateItem(
'dbs/' + self.test_database_default + '/colls/' +
self.test_collection_default, {
'data': 'data2',
'id': test_id2
}),
mock.call().CreateItem(
'dbs/' + self.test_database_default + '/colls/' +
self.test_collection_default, {
'data': 'data3',
'id': test_id3
})
]
logging.getLogger().info(returned_item)
cosmos_mock.assert_any_call(self.test_end_point,
{'masterKey': self.test_master_key})
cosmos_mock.assert_has_calls(expected_calls)
def test_upsert_document_default(self, cosmos_mock):
test_id = str(uuid.uuid4())
cosmos_mock.return_value.CreateItem.return_value = {'id': test_id}
self.cosmos = AzureCosmosDBHook(azure_cosmos_conn_id='azure_cosmos_test_key_id')
returned_item = self.cosmos.upsert_document({'id': test_id})
expected_calls = [mock.call().CreateItem(
'dbs/' + self.test_database_default + '/colls/' + self.test_collection_default,
{'id': test_id})]
cosmos_mock.assert_any_call(self.test_end_point, {'masterKey': self.test_master_key})
cosmos_mock.assert_has_calls(expected_calls)
logging.getLogger().info(returned_item)
self.assertEqual(returned_item['id'], test_id)
def test_insert_documents(self, cosmos_mock):
test_id1 = str(uuid.uuid4())
test_id2 = str(uuid.uuid4())
test_id3 = str(uuid.uuid4())
documents = [
{'id': test_id1, 'data': 'data1'},
{'id': test_id2, 'data': 'data2'},
{'id': test_id3, 'data': 'data3'}]
self.cosmos = AzureCosmosDBHook(azure_cosmos_conn_id='azure_cosmos_test_key_id')
returned_item = self.cosmos.insert_documents(documents)
expected_calls = [
mock.call().CreateItem(
'dbs/' + self.test_database_default + '/colls/' + self.test_collection_default,
{'data': 'data1', 'id': test_id1}),
mock.call().CreateItem(
'dbs/' + self.test_database_default + '/colls/' + self.test_collection_default,
{'data': 'data2', 'id': test_id2}),
mock.call().CreateItem(
'dbs/' + self.test_database_default + '/colls/' + self.test_collection_default,
{'data': 'data3', 'id': test_id3})]
logging.getLogger().info(returned_item)
cosmos_mock.assert_any_call(self.test_end_point, {'masterKey': self.test_master_key})
cosmos_mock.assert_has_calls(expected_calls)
def poke(self, context):
self.log.info("*** Intering poke")
hook = AzureCosmosDBHook(self.azure_cosmos_conn_id)
return hook.get_document(self.document_id, self.database_name, self.collection_name) is not None
def test_delete_database_exception(self, mock_cosmos):
hook = AzureCosmosDBHook(
azure_cosmos_conn_id='azure_cosmos_test_key_id')
self.assertRaises(AirflowException, hook.delete_database, None)
def test_create_container_exception(self, cosmos_mock):
self.cosmos = AzureCosmosDBHook(azure_cosmos_conn_id='azure_cosmos_test_key_id')
self.assertRaises(AirflowException, self.cosmos.create_collection, None)
def test_create_container_exception(self, mock_cosmos):
hook = AzureCosmosDBHook(
azure_cosmos_conn_id='azure_cosmos_test_key_id')
self.assertRaises(AirflowException, hook.create_collection, None)
def get_hook(self):
if self.conn_type == 'mysql':
from airflow.hooks.mysql_hook import MySqlHook
return MySqlHook(mysql_conn_id=self.conn_id)
elif self.conn_type == 'google_cloud_platform':
from airflow.gcp.hooks.bigquery import BigQueryHook
return BigQueryHook(bigquery_conn_id=self.conn_id)
elif self.conn_type == 'postgres':
from airflow.hooks.postgres_hook import PostgresHook
return PostgresHook(postgres_conn_id=self.conn_id)
elif self.conn_type == 'pig_cli':
from airflow.hooks.pig_hook import PigCliHook
return PigCliHook(pig_cli_conn_id=self.conn_id)
elif self.conn_type == 'hive_cli':
from airflow.hooks.hive_hooks import HiveCliHook
return HiveCliHook(hive_cli_conn_id=self.conn_id)
elif self.conn_type == 'presto':
from airflow.hooks.presto_hook import PrestoHook
return PrestoHook(presto_conn_id=self.conn_id)
elif self.conn_type == 'hiveserver2':
from airflow.hooks.hive_hooks import HiveServer2Hook
return HiveServer2Hook(hiveserver2_conn_id=self.conn_id)
elif self.conn_type == 'sqlite':
from airflow.hooks.sqlite_hook import SqliteHook
return SqliteHook(sqlite_conn_id=self.conn_id)
elif self.conn_type == 'jdbc':
from airflow.hooks.jdbc_hook import JdbcHook
return JdbcHook(jdbc_conn_id=self.conn_id)
elif self.conn_type == 'mssql':
from airflow.hooks.mssql_hook import MsSqlHook
return MsSqlHook(mssql_conn_id=self.conn_id)
elif self.conn_type == 'oracle':
from airflow.hooks.oracle_hook import OracleHook
return OracleHook(oracle_conn_id=self.conn_id)
elif self.conn_type == 'vertica':
from airflow.contrib.hooks.vertica_hook import VerticaHook
return VerticaHook(vertica_conn_id=self.conn_id)
elif self.conn_type == 'cloudant':
from airflow.contrib.hooks.cloudant_hook import CloudantHook
return CloudantHook(cloudant_conn_id=self.conn_id)
elif self.conn_type == 'jira':
from airflow.contrib.hooks.jira_hook import JiraHook
return JiraHook(jira_conn_id=self.conn_id)
elif self.conn_type == 'redis':
from airflow.contrib.hooks.redis_hook import RedisHook
return RedisHook(redis_conn_id=self.conn_id)
elif self.conn_type == 'wasb':
from airflow.contrib.hooks.wasb_hook import WasbHook
return WasbHook(wasb_conn_id=self.conn_id)
elif self.conn_type == 'docker':
from airflow.hooks.docker_hook import DockerHook
return DockerHook(docker_conn_id=self.conn_id)
elif self.conn_type == 'azure_data_lake':
from airflow.contrib.hooks.azure_data_lake_hook import AzureDataLakeHook
return AzureDataLakeHook(azure_data_lake_conn_id=self.conn_id)
elif self.conn_type == 'azure_cosmos':
from airflow.contrib.hooks.azure_cosmos_hook import AzureCosmosDBHook
return AzureCosmosDBHook(azure_cosmos_conn_id=self.conn_id)
elif self.conn_type == 'cassandra':
from airflow.contrib.hooks.cassandra_hook import CassandraHook
return CassandraHook(cassandra_conn_id=self.conn_id)
elif self.conn_type == 'mongo':
from airflow.contrib.hooks.mongo_hook import MongoHook
return MongoHook(conn_id=self.conn_id)
elif self.conn_type == 'gcpcloudsql':
from airflow.gcp.hooks.cloud_sql import CloudSqlDatabaseHook
return CloudSqlDatabaseHook(gcp_cloudsql_conn_id=self.conn_id)
elif self.conn_type == 'grpc':
from airflow.contrib.hooks.grpc_hook import GrpcHook
return GrpcHook(grpc_conn_id=self.conn_id)
raise AirflowException("Unknown hook type {}".format(self.conn_type))
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_create_database(self, cosmos_mock):
self.cosmos = AzureCosmosDBHook(
azure_cosmos_conn_id='azure_cosmos_test_key_id')
self.cosmos.create_database(self.test_database_name)
expected_calls = [
mock.call().CreateDatabase({'id': self.test_database_name})
]
cosmos_mock.assert_any_call(self.test_end_point,
{'masterKey': self.test_master_key})
cosmos_mock.assert_has_calls(expected_calls)
@mock.patch('azure.cosmos.cosmos_client.CosmosClient')
def test_create_database_exception(self, cosmos_mock):
self.cosmos = AzureCosmosDBHook(
azure_cosmos_conn_id='azure_cosmos_test_key_id')
self.assertRaises(AirflowException, self.cosmos.create_database, None)
@mock.patch('azure.cosmos.cosmos_client.CosmosClient')
def test_create_container_exception(self, cosmos_mock):
self.cosmos = AzureCosmosDBHook(
azure_cosmos_conn_id='azure_cosmos_test_key_id')
self.assertRaises(AirflowException, self.cosmos.create_collection,
None)
@mock.patch('azure.cosmos.cosmos_client.CosmosClient')
def test_create_container(self, cosmos_mock):
self.cosmos = AzureCosmosDBHook(
azure_cosmos_conn_id='azure_cosmos_test_key_id')
self.cosmos.create_collection(self.test_collection_name,
self.test_database_name)
expected_calls = [
mock.call().CreateContainer('dbs/test_database_name',
{'id': self.test_collection_name})
]
cosmos_mock.assert_any_call(self.test_end_point,
{'masterKey': self.test_master_key})
cosmos_mock.assert_has_calls(expected_calls)
@mock.patch('azure.cosmos.cosmos_client.CosmosClient')
def test_create_container_default(self, cosmos_mock):
self.cosmos = AzureCosmosDBHook(
azure_cosmos_conn_id='azure_cosmos_test_key_id')
self.cosmos.create_collection(self.test_collection_name)
expected_calls = [
mock.call().CreateContainer('dbs/test_database_default',
{'id': self.test_collection_name})
]
cosmos_mock.assert_any_call(self.test_end_point,
{'masterKey': self.test_master_key})
cosmos_mock.assert_has_calls(expected_calls)
@mock.patch('azure.cosmos.cosmos_client.CosmosClient')
def test_upsert_document_default(self, cosmos_mock):
test_id = str(uuid.uuid4())
cosmos_mock.return_value.CreateItem.return_value = {'id': test_id}
self.cosmos = AzureCosmosDBHook(
azure_cosmos_conn_id='azure_cosmos_test_key_id')
returned_item = self.cosmos.upsert_document({'id': test_id})
expected_calls = [
mock.call().CreateItem(
'dbs/' + self.test_database_default + '/colls/' +
self.test_collection_default, {'id': test_id})
]
cosmos_mock.assert_any_call(self.test_end_point,
{'masterKey': self.test_master_key})
cosmos_mock.assert_has_calls(expected_calls)
logging.getLogger().info(returned_item)
self.assertEqual(returned_item['id'], test_id)
@mock.patch('azure.cosmos.cosmos_client.CosmosClient')
def test_upsert_document(self, cosmos_mock):
test_id = str(uuid.uuid4())
cosmos_mock.return_value.CreateItem.return_value = {'id': test_id}
self.cosmos = AzureCosmosDBHook(
azure_cosmos_conn_id='azure_cosmos_test_key_id')
returned_item = self.cosmos.upsert_document(
{'data1': 'somedata'},
database_name=self.test_database_name,
collection_name=self.test_collection_name,
document_id=test_id)
expected_calls = [
mock.call().CreateItem(
'dbs/' + self.test_database_name + '/colls/' +
self.test_collection_name, {
'data1': 'somedata',
'id': test_id
})
]
cosmos_mock.assert_any_call(self.test_end_point,
{'masterKey': self.test_master_key})
cosmos_mock.assert_has_calls(expected_calls)
logging.getLogger().info(returned_item)
self.assertEqual(returned_item['id'], test_id)
@mock.patch('azure.cosmos.cosmos_client.CosmosClient')
def test_insert_documents(self, cosmos_mock):
test_id1 = str(uuid.uuid4())
test_id2 = str(uuid.uuid4())
test_id3 = str(uuid.uuid4())
documents = [{
'id': test_id1,
'data': 'data1'
}, {
'id': test_id2,
'data': 'data2'
}, {
'id': test_id3,
'data': 'data3'
}]
self.cosmos = AzureCosmosDBHook(
azure_cosmos_conn_id='azure_cosmos_test_key_id')
returned_item = self.cosmos.insert_documents(documents)
expected_calls = [
mock.call().CreateItem(
'dbs/' + self.test_database_default + '/colls/' +
self.test_collection_default, {
'data': 'data1',
'id': test_id1
}),
mock.call().CreateItem(
'dbs/' + self.test_database_default + '/colls/' +
self.test_collection_default, {
'data': 'data2',
'id': test_id2
}),
mock.call().CreateItem(
'dbs/' + self.test_database_default + '/colls/' +
self.test_collection_default, {
'data': 'data3',
'id': test_id3
})
]
logging.getLogger().info(returned_item)
cosmos_mock.assert_any_call(self.test_end_point,
{'masterKey': self.test_master_key})
cosmos_mock.assert_has_calls(expected_calls)
@mock.patch('azure.cosmos.cosmos_client.CosmosClient')
def test_delete_database(self, cosmos_mock):
self.cosmos = AzureCosmosDBHook(
azure_cosmos_conn_id='azure_cosmos_test_key_id')
self.cosmos.delete_database(self.test_database_name)
expected_calls = [mock.call().DeleteDatabase('dbs/test_database_name')]
cosmos_mock.assert_any_call(self.test_end_point,
{'masterKey': self.test_master_key})
cosmos_mock.assert_has_calls(expected_calls)
@mock.patch('azure.cosmos.cosmos_client.CosmosClient')
def test_delete_database_exception(self, cosmos_mock):
self.cosmos = AzureCosmosDBHook(
azure_cosmos_conn_id='azure_cosmos_test_key_id')
self.assertRaises(AirflowException, self.cosmos.delete_database, None)
@mock.patch('azure.cosmos.cosmos_client.CosmosClient')
def test_delete_container_exception(self, cosmos_mock):
self.cosmos = AzureCosmosDBHook(
azure_cosmos_conn_id='azure_cosmos_test_key_id')
self.assertRaises(AirflowException, self.cosmos.delete_collection,
None)
@mock.patch('azure.cosmos.cosmos_client.CosmosClient')
def test_delete_container(self, cosmos_mock):
self.cosmos = AzureCosmosDBHook(
azure_cosmos_conn_id='azure_cosmos_test_key_id')
self.cosmos.delete_collection(self.test_collection_name,
self.test_database_name)
expected_calls = [
mock.call().DeleteContainer(
'dbs/test_database_name/colls/test_collection_name')
]
cosmos_mock.assert_any_call(self.test_end_point,
{'masterKey': self.test_master_key})
cosmos_mock.assert_has_calls(expected_calls)
@mock.patch('azure.cosmos.cosmos_client.CosmosClient')
def test_delete_container_default(self, cosmos_mock):
self.cosmos = AzureCosmosDBHook(
azure_cosmos_conn_id='azure_cosmos_test_key_id')
self.cosmos.delete_collection(self.test_collection_name)
expected_calls = [
mock.call().DeleteContainer(
'dbs/test_database_default/colls/test_collection_name')
]
cosmos_mock.assert_any_call(self.test_end_point,
{'masterKey': self.test_master_key})
cosmos_mock.assert_has_calls(expected_calls)
def test_delete_database_exception(self, cosmos_mock):
self.cosmos = AzureCosmosDBHook(
azure_cosmos_conn_id='azure_cosmos_test_key_id')
self.assertRaises(AirflowException, self.cosmos.delete_database, None)
def test_delete_database(self, cosmos_mock):
self.cosmos = AzureCosmosDBHook(azure_cosmos_conn_id='azure_cosmos_test_key_id')
self.cosmos.delete_database(self.test_database_name)
expected_calls = [mock.call().DeleteDatabase('dbs/test_database_name')]
cosmos_mock.assert_any_call(self.test_end_point, {'masterKey': self.test_master_key})
cosmos_mock.assert_has_calls(expected_calls)
def test_create_container_exception(self, cosmos_mock):
self.cosmos = AzureCosmosDBHook(
azure_cosmos_conn_id='azure_cosmos_test_key_id')
self.assertRaises(AirflowException, self.cosmos.create_collection,
None)
def test_delete_database_exception(self, cosmos_mock):
self.cosmos = AzureCosmosDBHook(azure_cosmos_conn_id='azure_cosmos_test_key_id')
self.assertRaises(AirflowException, self.cosmos.delete_database, None)
def poke(self, context):
self.log.info("*** Intering poke")
hook = AzureCosmosDBHook(self.azure_cosmos_conn_id)
return hook.get_document(self.document_id, self.database_name,
self.collection_name) is not None
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_create_database(self, cosmos_mock):
self.cosmos = AzureCosmosDBHook(azure_cosmos_conn_id='azure_cosmos_test_key_id')
self.cosmos.create_database(self.test_database_name)
expected_calls = [mock.call().CreateDatabase({'id': self.test_database_name})]
cosmos_mock.assert_any_call(self.test_end_point, {'masterKey': self.test_master_key})
cosmos_mock.assert_has_calls(expected_calls)
@mock.patch('azure.cosmos.cosmos_client.CosmosClient')
def test_create_database_exception(self, cosmos_mock):
self.cosmos = AzureCosmosDBHook(azure_cosmos_conn_id='azure_cosmos_test_key_id')
self.assertRaises(AirflowException, self.cosmos.create_database, None)
@mock.patch('azure.cosmos.cosmos_client.CosmosClient')
def test_create_container_exception(self, cosmos_mock):
self.cosmos = AzureCosmosDBHook(azure_cosmos_conn_id='azure_cosmos_test_key_id')
self.assertRaises(AirflowException, self.cosmos.create_collection, None)
@mock.patch('azure.cosmos.cosmos_client.CosmosClient')
def test_create_container(self, cosmos_mock):
self.cosmos = AzureCosmosDBHook(azure_cosmos_conn_id='azure_cosmos_test_key_id')
self.cosmos.create_collection(self.test_collection_name, self.test_database_name)
expected_calls = [mock.call().CreateContainer(
'dbs/test_database_name',
{'id': self.test_collection_name})]
cosmos_mock.assert_any_call(self.test_end_point, {'masterKey': self.test_master_key})
cosmos_mock.assert_has_calls(expected_calls)
@mock.patch('azure.cosmos.cosmos_client.CosmosClient')
def test_create_container_default(self, cosmos_mock):
self.cosmos = AzureCosmosDBHook(azure_cosmos_conn_id='azure_cosmos_test_key_id')
self.cosmos.create_collection(self.test_collection_name)
expected_calls = [mock.call().CreateContainer(
'dbs/test_database_default',
{'id': self.test_collection_name})]
cosmos_mock.assert_any_call(self.test_end_point, {'masterKey': self.test_master_key})
cosmos_mock.assert_has_calls(expected_calls)
@mock.patch('azure.cosmos.cosmos_client.CosmosClient')
def test_upsert_document_default(self, cosmos_mock):
test_id = str(uuid.uuid4())
cosmos_mock.return_value.CreateItem.return_value = {'id': test_id}
self.cosmos = AzureCosmosDBHook(azure_cosmos_conn_id='azure_cosmos_test_key_id')
returned_item = self.cosmos.upsert_document({'id': test_id})
expected_calls = [mock.call().CreateItem(
'dbs/' + self.test_database_default + '/colls/' + self.test_collection_default,
{'id': test_id})]
cosmos_mock.assert_any_call(self.test_end_point, {'masterKey': self.test_master_key})
cosmos_mock.assert_has_calls(expected_calls)
logging.getLogger().info(returned_item)
self.assertEqual(returned_item['id'], test_id)
@mock.patch('azure.cosmos.cosmos_client.CosmosClient')
def test_upsert_document(self, cosmos_mock):
test_id = str(uuid.uuid4())
cosmos_mock.return_value.CreateItem.return_value = {'id': test_id}
self.cosmos = AzureCosmosDBHook(azure_cosmos_conn_id='azure_cosmos_test_key_id')
returned_item = self.cosmos.upsert_document(
{'data1': 'somedata'},
database_name=self.test_database_name,
collection_name=self.test_collection_name,
document_id=test_id)
expected_calls = [mock.call().CreateItem(
'dbs/' + self.test_database_name + '/colls/' + self.test_collection_name,
{'data1': 'somedata', 'id': test_id})]
cosmos_mock.assert_any_call(self.test_end_point, {'masterKey': self.test_master_key})
cosmos_mock.assert_has_calls(expected_calls)
logging.getLogger().info(returned_item)
self.assertEqual(returned_item['id'], test_id)
@mock.patch('azure.cosmos.cosmos_client.CosmosClient')
def test_insert_documents(self, cosmos_mock):
test_id1 = str(uuid.uuid4())
test_id2 = str(uuid.uuid4())
test_id3 = str(uuid.uuid4())
documents = [
{'id': test_id1, 'data': 'data1'},
{'id': test_id2, 'data': 'data2'},
{'id': test_id3, 'data': 'data3'}]
self.cosmos = AzureCosmosDBHook(azure_cosmos_conn_id='azure_cosmos_test_key_id')
returned_item = self.cosmos.insert_documents(documents)
expected_calls = [
mock.call().CreateItem(
'dbs/' + self.test_database_default + '/colls/' + self.test_collection_default,
{'data': 'data1', 'id': test_id1}),
mock.call().CreateItem(
'dbs/' + self.test_database_default + '/colls/' + self.test_collection_default,
{'data': 'data2', 'id': test_id2}),
mock.call().CreateItem(
'dbs/' + self.test_database_default + '/colls/' + self.test_collection_default,
{'data': 'data3', 'id': test_id3})]
logging.getLogger().info(returned_item)
cosmos_mock.assert_any_call(self.test_end_point, {'masterKey': self.test_master_key})
cosmos_mock.assert_has_calls(expected_calls)
@mock.patch('azure.cosmos.cosmos_client.CosmosClient')
def test_delete_database(self, cosmos_mock):
self.cosmos = AzureCosmosDBHook(azure_cosmos_conn_id='azure_cosmos_test_key_id')
self.cosmos.delete_database(self.test_database_name)
expected_calls = [mock.call().DeleteDatabase('dbs/test_database_name')]
cosmos_mock.assert_any_call(self.test_end_point, {'masterKey': self.test_master_key})
cosmos_mock.assert_has_calls(expected_calls)
@mock.patch('azure.cosmos.cosmos_client.CosmosClient')
def test_delete_database_exception(self, cosmos_mock):
self.cosmos = AzureCosmosDBHook(azure_cosmos_conn_id='azure_cosmos_test_key_id')
self.assertRaises(AirflowException, self.cosmos.delete_database, None)
@mock.patch('azure.cosmos.cosmos_client.CosmosClient')
def test_delete_container_exception(self, cosmos_mock):
self.cosmos = AzureCosmosDBHook(azure_cosmos_conn_id='azure_cosmos_test_key_id')
self.assertRaises(AirflowException, self.cosmos.delete_collection, None)
@mock.patch('azure.cosmos.cosmos_client.CosmosClient')
def test_delete_container(self, cosmos_mock):
self.cosmos = AzureCosmosDBHook(azure_cosmos_conn_id='azure_cosmos_test_key_id')
self.cosmos.delete_collection(self.test_collection_name, self.test_database_name)
expected_calls = [mock.call().DeleteContainer('dbs/test_database_name/colls/test_collection_name')]
cosmos_mock.assert_any_call(self.test_end_point, {'masterKey': self.test_master_key})
cosmos_mock.assert_has_calls(expected_calls)
@mock.patch('azure.cosmos.cosmos_client.CosmosClient')
def test_delete_container_default(self, cosmos_mock):
self.cosmos = AzureCosmosDBHook(azure_cosmos_conn_id='azure_cosmos_test_key_id')
self.cosmos.delete_collection(self.test_collection_name)
expected_calls = [mock.call().DeleteContainer('dbs/test_database_default/colls/test_collection_name')]
cosmos_mock.assert_any_call(self.test_end_point, {'masterKey': self.test_master_key})
cosmos_mock.assert_has_calls(expected_calls)
