Add support for tags in search_entities and update tests

Author: tomchop

tests/api.py

@@ -66,13 +66,17 @@ def test_search_entities(self, mock_post):
         mock_post.return_value = mock_response
 
         result = self.api.search_entities(
-            name="test_entity", description="test_description"
+            name="test_entity", description="test_description", tags=["tag1"]
         )
         self.assertEqual(result, [{"name": "test_entity"}])
         mock_post.assert_called_with(
             "http://fake-url/api/v2/entities/search",
             json={
-                "query": {"name": "test_entity", "description": "test_description"},
+                "query": {
+                    "name": "test_entity",
+                    "description": "test_description",
+                    "tags": ["tag1"],
+                },
                 "count": 100,
                 "page": 0,
             },

tests/e2e.py

@@ -5,6 +5,10 @@
 from yeti import errors
 from yeti.api import YetiApi
 
+os.environ["YETI_ENDPOINT"] = "http://localhost:3000/"
+YETI_API_KEY = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJuYW1lIjoibWNwIiwic3ViIjoieWV0aSIsInNjb3BlcyI6WyJhbGwiXSwiY3JlYXRlZCI6IjIwMjUtMDYtMDJUMTk6NDI6MjcuNDcxNzA5WiIsImV4cCI6bnVsbCwibGFzdF91c2VkIjpudWxsLCJlbmFibGVkIjp0cnVlLCJleHBpcmVkIjpmYWxzZX0.qw6TlHkkxol--qjdOdnporlyYP0i2uJGgjzdKybpX2U"
+os.environ["YETI_API_KEY"] = YETI_API_KEY
+
 
 class YetiEndToEndTest(unittest.TestCase):
     def setUp(self):
@@ -60,6 +64,23 @@ def test_search_entities(self):
         self.assertEqual(result[0]["name"], "testSearch")
         self.assertEqual(result[0]["tags"][0]["name"], "testtag")
 
+    def test_search_entities_with_tags(self):
+        self.api.auth_api_key(os.getenv("YETI_API_KEY"))
+        self.api.new_entity(
+            {
+                "name": "testSearchWithTags",
+                "type": "malware",
+                "description": "test",
+            },
+            tags=["testtag1", "testtag2"],
+        )
+        time.sleep(5)
+        result = self.api.search_entities(
+            name="testSear", description="tes", tags=["testtag1"]
+        )
+        self.assertEqual(len(result), 1, result)
+        self.assertEqual(result[0]["name"], "testSearchWithTags")
+
     def test_get_multiple_entities(self):
         self.api.auth_api_key(os.getenv("YETI_API_KEY"))
         self.api.new_entity(

yeti/api.py

@@ -224,6 +224,8 @@ def search_indicators(
           pattern: The pattern of the indicator to search for.
           description: The description of the indicator to search for. (substring match)
           tags: The tags of the indicator to search for.
+          count: The number of results to return (default is 100, which means all).
+          page: The page of results to return (default is 0, which means the first page).
 
         Returns:
           The response from the API; a list of dicts representing indicators.
@@ -299,6 +301,7 @@ def search_entities(
         name: str | None = None,
         entity_type: str | None = None,
         description: str | None = None,
+        tags: list[str] | None = None,
         count: int = 100,
         page: int = 0,
     ) -> list[YetiObject]:
@@ -310,6 +313,7 @@ def search_entities(
             name: The name of the entity to search for (substring match).
             entity_type: The type of the entity to search for.
             description: The description of the entity to search for. (substring match)
+            tags: The tags of the entity to search for.
             count: The number of results to return (default is 100, which means all).
             page: The page of results to return (default is 0, which means the first page).
 
@@ -326,6 +330,8 @@ def search_entities(
             query["type"] = entity_type
         if description:
             query["description"] = description
+        if tags:
+            query["tags"] = tags
 
         params = {"query": query, "count": count, "page": page}
         response = self.do_request(