@@ -610,7 +610,7 @@ def test_assert_valid_hook_call(self, mock_hook) -> None:
610
610
class TestCloudDataCatalogListTagsOperator (TestCase ):
611
611
@mock .patch (
612
612
"airflow.providers.google.cloud.operators.datacatalog.CloudDataCatalogHook" ,
613
- ** { " return_value. list_tags. return_value" : [TEST_TAG ]}, # type: ignore
613
+ return_value = mock . MagicMock ( list_tags = mock . MagicMock ( return_value = [TEST_TAG ])),
614
614
)
615
615
def test_assert_valid_hook_call (self , mock_hook ) -> None :
616
616
task = CloudDataCatalogListTagsOperator (
@@ -777,7 +777,7 @@ class TestCloudDataCatalogUpdateTagOperator(TestCase):
777
777
def test_assert_valid_hook_call (self , mock_hook ) -> None :
778
778
task = CloudDataCatalogUpdateTagOperator (
779
779
task_id = "task_id" ,
780
- tag = TEST_TAG_ID ,
780
+ tag = Tag ( name = TEST_TAG_ID ) ,
781
781
update_mask = TEST_UPDATE_MASK ,
782
782
location = TEST_LOCATION ,
783
783
entry_group = TEST_ENTRY_GROUP_ID ,
@@ -796,7 +796,7 @@ def test_assert_valid_hook_call(self, mock_hook) -> None:
796
796
impersonation_chain = TEST_IMPERSONATION_CHAIN ,
797
797
)
798
798
mock_hook .return_value .update_tag .assert_called_once_with (
799
- tag = TEST_TAG_ID ,
799
+ tag = Tag ( name = TEST_TAG_ID ) ,
800
800
update_mask = TEST_UPDATE_MASK ,
801
801
location = TEST_LOCATION ,
802
802
entry_group = TEST_ENTRY_GROUP_ID ,
@@ -814,7 +814,7 @@ class TestCloudDataCatalogUpdateTagTemplateOperator(TestCase):
814
814
def test_assert_valid_hook_call (self , mock_hook ) -> None :
815
815
task = CloudDataCatalogUpdateTagTemplateOperator (
816
816
task_id = "task_id" ,
817
- tag_template = TEST_TAG_TEMPLATE_ID ,
817
+ tag_template = TagTemplate ( name = TEST_TAG_TEMPLATE_ID ) ,
818
818
update_mask = TEST_UPDATE_MASK ,
819
819
location = TEST_LOCATION ,
820
820
tag_template_id = TEST_TAG_TEMPLATE_ID ,
@@ -831,7 +831,7 @@ def test_assert_valid_hook_call(self, mock_hook) -> None:
831
831
impersonation_chain = TEST_IMPERSONATION_CHAIN ,
832
832
)
833
833
mock_hook .return_value .update_tag_template .assert_called_once_with (
834
- tag_template = TEST_TAG_TEMPLATE_ID ,
834
+ tag_template = TagTemplate ( name = TEST_TAG_TEMPLATE_ID ) ,
835
835
update_mask = TEST_UPDATE_MASK ,
836
836
location = TEST_LOCATION ,
837
837
tag_template_id = TEST_TAG_TEMPLATE_ID ,
0 commit comments