Refine identity protection unit testing

CrowdStrike · Jul 20, 2023 · aba7f14 · aba7f14
1 parent ca31e11
commit aba7f14
Showing 1 changed file with 24 additions and 23 deletions.
diff --git a/tests/test_identity_protection.py b/tests/test_identity_protection.py
@@ -18,33 +18,38 @@
 falcon = IdentityProtection(auth_object=config)
 AllowedResponses = [200, 429]
 
+TEST_QUERY = """
+query ($after: Cursor) {
+  entities(types: [USER], archived: false, learned: false, first: 5, after: $after) {
+    nodes {
+      primaryDisplayName
+      secondaryDisplayName
+      accounts {
+        ... on ActiveDirectoryAccountDescriptor {
+          domain
+        }
+      }
+    }
+    pageInfo {
+      hasNextPage
+      endCursor
+    }
+  }
+}
+"""
 
 class TestIdentityProtection:
     def idp_graphql(self):
         payload = {"query":"{\n  entities(first: 1)\n  {\n    nodes {\n      entityId    \n    }\n  }\n}"}
-        result = falcon.GraphQL(body=payload)
+        result = falcon.GraphQL(query=TEST_QUERY, variables={"after", ""})
         if not isinstance(result, dict):
             result = json.loads(result.decode())
         else:
             result = result["body"]
-        if "extensions" in result:
-            if result["extensions"]["remainingPoints"] > 0:
-                return True
-            else:
-                pytest.skip("Identity protection API failure")
-                # return False
-        else:
-            # Prolly failed login, check yer API key
-            # return False
-            pytest.skip("Identity protection API failure")
+        if result.get("data", {}).get("entities", {}).get("pageInfo", {}).get("hasNextPage", None):
+            next_page = result["data"].get("entities", {}).get("pageInfo", {}).get("endCursor", None)
+            result = falcon.graphql(query=TEST_QUERY, variables={"after": next_page})["body"]
 
-    def idp_graphql_keywords(self):
-        test_query = "{\n  entities(first: 1)\n  {\n    nodes {\n      entityId    \n    }\n  }\n}"
-        result = falcon.graphql(query=test_query, variables={"after": "whatever"})
-        if not isinstance(result, dict):
-            result = json.loads(result.decode())
-        else:
-            result = result["body"]
         if "extensions" in result:
             if result["extensions"]["remainingPoints"] > 0:
                 return True
@@ -56,14 +61,10 @@ def idp_graphql_keywords(self):
             # return False
             pytest.skip("Identity protection API failure")
 
+
     @pytest.mark.skipif(falcon.base_url.lower() in ["https://api.laggar.gcw.crowdstrike.com","usgov1"],
                         reason="Unit testing unavailable on US-GOV-1"
                         )
     def test_graphql(self):
         assert self.idp_graphql() is True
 
-    @pytest.mark.skipif(falcon.base_url.lower() in ["https://api.laggar.gcw.crowdstrike.com","usgov1"],
-                        reason="Unit testing unavailable on US-GOV-1"
-                        )
-    def test_graphql_keywords(self):
-        assert self.idp_graphql_keywords() is True