1#
2# Licensed to the Apache Software Foundation (ASF) under one
3# or more contributor license agreements. See the NOTICE file
4# distributed with this work for additional information
5# regarding copyright ownership. The ASF licenses this file
6# to you under the Apache License, Version 2.0 (the
7# "License"); you may not use this file except in compliance
8# with the License. You may obtain a copy of the License at
9#
10# http://www.apache.org/licenses/LICENSE-2.0
11#
12# Unless required by applicable law or agreed to in writing,
13# software distributed under the License is distributed on an
14# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15# KIND, either express or implied. See the License for the
16# specific language governing permissions and limitations
17# under the License.
18from __future__ import annotations
19
20import uuid
21from unittest import mock
22
23import pytest
24
25from airflow.models import Connection
26from airflow.providers.microsoft.azure.operators.cosmos import AzureCosmosInsertDocumentOperator
27
28
29class TestAzureCosmosDbHook:
30 # Set up an environment to test with
31 @pytest.fixture(autouse=True)
32 def setup_test_cases(self, create_mock_connection):
33 # set up some test variables
34 self.test_end_point = "https://test_endpoint:443"
35 self.test_master_key = "magic_test_key"
36 self.test_database_name = "test_database_name"
37 self.test_collection_name = "test_collection_name"
38 self.test_partition_key = "test_partition_key"
39 create_mock_connection(
40 Connection(
41 conn_id="azure_cosmos_test_key_id",
42 conn_type="azure_cosmos",
43 login=self.test_end_point,
44 password=self.test_master_key,
45 extra={
46 "database_name": self.test_database_name,
47 "collection_name": self.test_collection_name,
48 "partition_key": self.test_partition_key,
49 },
50 )
51 )
52
53 @mock.patch("airflow.providers.microsoft.azure.hooks.cosmos.CosmosClient")
54 def test_insert_document(self, cosmos_mock):
55 test_id = str(uuid.uuid4())
56 cosmos_mock.return_value.CreateItem.return_value = {"id": test_id}
57 op = AzureCosmosInsertDocumentOperator(
58 database_name=self.test_database_name,
59 collection_name=self.test_collection_name,
60 document={"id": test_id, "data": "sometestdata"},
61 azure_cosmos_conn_id="azure_cosmos_test_key_id",
62 task_id="azure_cosmos_sensor",
63 )
64
65 expected_calls = [
66 mock.call()
67 .get_database_client("test_database_name")
68 .get_container_client("test_collection_name")
69 .upsert_item({"data": "sometestdata", "id": test_id})
70 ]
71
72 op.execute(None)
73 cosmos_mock.assert_any_call(self.test_end_point, {"masterKey": self.test_master_key})
74 cosmos_mock.assert_has_calls(expected_calls)