本文整理汇总了Python中airflow.contrib.hooks.azure_cosmos_hook.AzureCosmosDBHook类的典型用法代码示例。如果您正苦于以下问题:Python AzureCosmosDBHook类的具体用法?Python AzureCosmosDBHook怎么用?Python AzureCosmosDBHook使用的例子?那么, 这里精选的类代码示例或许可以为您提供帮助。
在下文中一共展示了AzureCosmosDBHook类的9个代码示例,这些例子默认根据受欢迎程度排序。您可以为喜欢或者感觉有用的代码点赞,您的评价将有助于系统推荐出更棒的Python代码示例。
示例1: test_create_container_default
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)
示例2: execute
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)
示例3: test_upsert_document_default
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)
示例4: test_insert_documents
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)
示例5: test_create_container_exception
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)
示例6: TestAzureCosmosDbHook
class TestAzureCosmosDbHook(unittest.TestCase):
# Set up an environment to test with
def setUp(self):
# set up some test variables
self.test_end_point = 'https://test_endpoint:443'
self.test_master_key = 'magic_test_key'
self.test_database_name = 'test_database_name'
self.test_collection_name = 'test_collection_name'
self.test_database_default = 'test_database_default'
self.test_collection_default = 'test_collection_default'
configuration.load_test_config()
db.merge_conn(
Connection(
conn_id='azure_cosmos_test_key_id',
conn_type='azure_cosmos',
login=self.test_end_point,
password=self.test_master_key,
extra=json.dumps({'database_name': self.test_database_default,
'collection_name': self.test_collection_default})
)
)
@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())
#.........这里部分代码省略.........
示例7: test_delete_database_exception
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)
示例8: test_delete_database
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)
示例9: poke
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