test_azure_cosmos_9.py 1.1 KB

12345678910111213141516171819202122232425262728293031
  1. def test_upsert_document(self, mock_cosmos):
  2. test_id = str(uuid.uuid4())
  3. # fmt: off
  4. (mock_cosmos
  5. .return_value
  6. .get_database_client
  7. .return_value
  8. .get_container_client
  9. .return_value
  10. .upsert_item
  11. .return_value) = {'id': test_id}
  12. # fmt: on
  13. hook = AzureCosmosDBHook(azure_cosmos_conn_id='azure_cosmos_test_key_id')
  14. returned_item = hook.upsert_document(
  15. {'data1': 'somedata'},
  16. database_name=self.test_database_name,
  17. collection_name=self.test_collection_name,
  18. document_id=test_id,
  19. )
  20. expected_calls = [
  21. mock.call()
  22. .get_database_client('test_database_name')
  23. .get_container_client('test_collection_name')
  24. .upsert_item({'data1': 'somedata', 'id': test_id})
  25. ]
  26. mock_cosmos.assert_any_call(self.test_end_point, {'masterKey': self.test_master_key})
  27. mock_cosmos.assert_has_calls(expected_calls)
  28. logging.getLogger().info(returned_item)
  29. assert returned_item['id'] == test_id