Update unti tests + formatting

Author: marc-sift

python/lib/sift_py/rule/_service_test.py

@@ -93,6 +93,7 @@ def test_rule_service_load_rules_from_yaml(rule_service):
         "assignee": "assignee@abc.com",
         "type": "review",
         "asset_names": ["asset"],
+        "tag_names": ["tag1"],
     }
     with mock.patch.object(RuleService, "create_or_update_rule"):
         with mock.patch(
@@ -115,6 +116,7 @@ def test_rule_service_load_rules_from_yaml(rule_service):
             assert rule_config.expression == rule_yaml["expression"]
             assert rule_config.action.assignee == rule_yaml["assignee"]
             assert rule_config.asset_names == rule_yaml["asset_names"]
+            assert rule_config.tag_names == rule_yaml["tag1"]
             assert isinstance(rule_config.action, RuleActionCreateDataReviewAnnotation)
 
 

python/lib/sift_py/rule/config.py

@@ -135,6 +135,8 @@ def interpolate_sub_expressions(
 
 
 class RuleAction(ABC):
+    tags: Optional[List[str]]
+
     @abstractmethod
     def kind(self) -> RuleActionKind:
         pass

python/lib/sift_py/rule/service.py

@@ -414,10 +414,14 @@ def _update_req_from_rule_config(
             if config.tag_names
             else None
         )
-        annotation_tags = list_tags_impl(
-            self._tag_service_stub,
-            names=[tag for tag in config.action.tags],
-            tag_type=TagType.TAG_TYPE_ANNOTATION,
+        annotation_tags = (
+            list_tags_impl(
+                self._tag_service_stub,
+                names=[tag for tag in config.action.tags],
+                tag_type=TagType.TAG_TYPE_ANNOTATION,
+            )
+            if config.action.tags
+            else None
         )
 
         actions = []
@@ -427,6 +431,10 @@ def _update_req_from_rule_config(
                 "Please contact the Sift team for assistance."
             )
         elif config.action.kind() == RuleActionKind.ANNOTATION:
+            annotation_tag_ids = (
+                [tag.tag_id for tag in annotation_tags] if annotation_tags else None
+            )
+
             if isinstance(config.action, RuleActionCreateDataReviewAnnotation):
                 assignee = config.action.assignee
                 user_id = None
@@ -446,7 +454,7 @@ def _update_req_from_rule_config(
                         annotation=AnnotationActionConfiguration(
                             assigned_to_user_id=user_id,
                             annotation_type=AnnotationType.ANNOTATION_TYPE_DATA_REVIEW,
-                            tag_ids=annotation_tags,
+                            tag_ids=annotation_tag_ids,
                         )
                     ),
                 )
@@ -457,7 +465,7 @@ def _update_req_from_rule_config(
                     configuration=RuleActionConfiguration(
                         annotation=AnnotationActionConfiguration(
                             annotation_type=AnnotationType.ANNOTATION_TYPE_PHASE,
-                            tag_ids=annotation_tags,
+                            tag_ids=annotation_tag_ids,
                         )
                     ),
                 )
@@ -640,7 +648,10 @@ def _get_assets(self, names: List[str] = [], ids: List[str] = []) -> List[Asset]
             return list_assets_impl(self._asset_service_stub, names, ids)
 
     def _get_tags(
-        self, names: List[str] = [], ids: List[str] = [], tag_type: Optional[TagType] = None
+        self,
+        names: List[str] = [],
+        ids: List[str] = [],
+        tag_type: TagType.ValueType = TagType.TAG_TYPE_UNSPECIFIED,
     ) -> List[Tag]:
         if self._enable_caching:
             return self._get_tags_cached(tuple(sorted(names)), tuple(sorted(ids)), tag_type)
@@ -665,8 +676,11 @@ def _get_assets_cached(self, names: Tuple[str], ids: Tuple[str]) -> List[Asset]:
 
     @cache
     def _get_tags_cached(
-        self, names: Tuple[str], ids: Tuple[str], tag_type: Optional[TagType] = None
-    ) -> List[Asset]:
+        self,
+        names: Tuple[str],
+        ids: Tuple[str],
+        tag_type: TagType.ValueType = TagType.TAG_TYPE_UNSPECIFIED,
+    ) -> List[Tag]:
         return list_tags_impl(self._tag_service_stub, names, ids, tag_type)
 
     @cache

python/lib/sift_py/tag/_internal/shared.py

@@ -9,7 +9,7 @@ def list_tags_impl(
     tag_service_stub: TagServiceStub,
     names: Optional[Union[Tuple[str], List[str]]] = None,
     ids: Optional[Union[Tuple[str], List[str]]] = None,
-    tag_type: Optional[TagType] = None,
+    tag_type: TagType.ValueType = TagType.TAG_TYPE_UNSPECIFIED,
 ) -> List[Tag]:
     """
     Lists tags in an organization.
@@ -25,26 +25,27 @@ def list_tags_impl(
     """
 
     def get_tags_with_filter(
-        tag_service_stub: TagServiceStub, cel_filter: str, tag_type: Optional[TagType]
+        tag_service_stub: TagServiceStub,
+        cel_filter: str,
+        tag_type: TagType.ValueType,
     ) -> List[Tag]:
         tags: List[Tag] = []
         next_page_token = ""
         while True:
-            req_kwargs = {
-                "filter": cel_filter,
-                "page_size": 1_000,
-                "page_token": next_page_token,
-            }
-            if tag_type is not None:
-                req_kwargs["tag_type"] = tag_type
-            req = ListTagsRequest(**req_kwargs)
+            req = ListTagsRequest(
+                filter=cel_filter,
+                page_size=1_000,
+                page_token=next_page_token,
+                tag_type=tag_type,
+            )
             res = cast(ListTagsResponse, tag_service_stub.ListTags(req))
             tags.extend(res.tags)
 
             if not res.next_page_token:
                 break
             next_page_token = res.next_page_token
 
+        print(tags)
         return tags
 
     if names is None:
@@ -55,13 +56,15 @@ def get_tags_with_filter(
     results: List[Tag] = []
     if names:
         names_cel = cel_in("name", names)
-        results.append(get_tags_with_filter(tag_service_stub, names_cel, tag_type))
+        results.extend(get_tags_with_filter(tag_service_stub, names_cel, tag_type))
     if ids:
         ids_cel = cel_in("tag_id", ids)
-        results.append(get_tags_with_filter(tag_service_stub, ids_cel, tag_type))
+        results.extend(get_tags_with_filter(tag_service_stub, ids_cel, tag_type))
     if not names and not ids:
         # If no filter, but tag_type is specified, fetch all tags of that type
         if tag_type is not None:
-            results.append(get_tags_with_filter(tag_service_stub, "", tag_type))
+            results.extend(get_tags_with_filter(tag_service_stub, "", tag_type))
+
+    print(names, ids, results)
 
     return results