Merge pull request #719 from Ostorlab/fix/agents-pagination

Fix agents pagination

src/ostorlab/serve_app/types.py

@@ -709,7 +709,7 @@ def resolve_args(
             return AgentArgumentsType(args=args)
 
 
-class AgentsType(graphene.ObjectType):
+class OxoAgentsType(graphene.ObjectType):
     """Graphene object type for a list of agents."""
 
     agents = graphene.List(AgentType, required=True)
@@ -724,7 +724,7 @@ class AgentGroupType(graphene_sqlalchemy.SQLAlchemyObjectType):
 
     key = graphene.String()
     agents = graphene.Field(
-        AgentsType,
+        OxoAgentsType,
         required=True,
         page=graphene.Int(required=False),
         number_elements=graphene.Int(required=False),
@@ -762,16 +762,16 @@ def resolve_agents(
         info: graphql_base.ResolveInfo,
         page: int = None,
         number_elements: int = DEFAULT_NUMBER_ELEMENTS,
-    ) -> AgentsType:
+    ) -> OxoAgentsType:
         """Resolve agents query.
         Args:
-            self (models.AgentGroup): The agent group object.
-            info (graphql_base.ResolveInfo): GraphQL resolve info.
+            self: The agent group object.
+            info: GraphQL resolve info.
         Returns:
-            AgentsType: List of agents.
+            List of agents.
         """
         if number_elements <= 0:
-            return AgentsType(agents=[])
+            return OxoAgentsType(agents=[])
 
         with models.Database() as session:
             agents = (
@@ -789,9 +789,9 @@ def resolve_agents(
                     has_next=page.has_next(),
                     has_previous=page.has_previous(),
                 )
-                return OxoVulnerabilitiesType(agents=page, page_info=page_info)
+                return OxoAgentsType(agents=page, page_info=page_info)
             else:
-                return AgentsType(agents=agents)
+                return OxoAgentsType(agents=agents)
 
     def resolve_asset_types(
         self: models.AgentGroup, info: graphql_base.ResolveInfo

tests/conftest.py

@@ -926,6 +926,41 @@ def agent_group(
         return agent_group
 
 
+@pytest.fixture
+def agent_group_multiple_agents(
+    clean_db: None, mocker: plugin.MockerFixture, db_engine_path: str
+) -> models.AgentGroup:
+    """Create dummy agent group."""
+    mocker.patch.object(models, "ENGINE_URL", db_engine_path)
+    with models.Database() as session:
+        agent1 = models.Agent.create(
+            key="agent/ostorlab/agent1",
+        )
+        agent2 = models.Agent.create(
+            key="agent/ostorlab/agent2",
+        )
+        agent3 = models.Agent.create(
+            key="agent/ostorlab/agent3",
+        )
+        agent_group = models.AgentGroup(
+            name="Agent Group 1",
+            description="Agent Group 1 description",
+            created_time=datetime.datetime.now(),
+        )
+        session.add(agent_group)
+        session.commit()
+        models.AgentGroupMapping.create(
+            agent_group_id=agent_group.id, agent_id=agent1.id
+        )
+        models.AgentGroupMapping.create(
+            agent_group_id=agent_group.id, agent_id=agent2.id
+        )
+        models.AgentGroupMapping.create(
+            agent_group_id=agent_group.id, agent_id=agent3.id
+        )
+        return agent_group
+
+
 @pytest.fixture
 def multiple_assets_scan(
     mocker: plugin.MockerFixture, db_engine_path: str

tests/serve_app/oxo_test.py

@@ -3211,3 +3211,51 @@ def testPublishAgentGroup_withoutNameAndAgentArgs_shouldPersistAgentGroup(
     assert len(ag["agents"]["agents"]) == 1
     assert ag["agents"]["agents"][0]["key"] == "agent_key"
     assert len(ag["agents"]["agents"][0]["args"]["args"]) == 0
+
+
+def testQueryAgentGroup_withAgentPagination_shouldReturnPaginatedListOfAgents(
+    authenticated_flask_client: testing.FlaskClient,
+    agent_group_multiple_agents: models.AgentGroup,
+    mocker: plugin.MockerFixture,
+    db_engine_path: str,
+) -> None:
+    """Ensure the query agent group returns a paginated list of agents."""
+    mocker.patch.object(models, "ENGINE_URL", db_engine_path)
+    query = """
+        query QueryAgentGroup ($agentGroupIds: [Int!], $page: Int, $numberElements: Int){
+            agentGroups (agentGroupIds: $agentGroupIds) {
+            agentGroups {
+                agents (page: $page, numberElements: $numberElements) {
+                    agents {
+                        key
+                    }
+                    pageInfo{
+                      count
+                      numPages
+                      hasNext
+                      hasPrevious
+                    }
+                }
+            }
+            }
+        }
+    """
+    variables = {
+        "agentGroupIds": [agent_group_multiple_agents.id],
+        "page": 1,
+        "numberElements": 2,
+    }
+
+    response = authenticated_flask_client.post(
+        "/graphql", json={"query": query, "variables": variables}
+    )
+
+    assert response.status_code == 200, response.get_json()
+    agents = response.get_json()["data"]["agentGroups"]["agentGroups"][0]["agents"]
+    assert len(agents["agents"]) == 2
+    assert agents["agents"][0]["key"] == "agent/ostorlab/agent1"
+    assert agents["agents"][1]["key"] == "agent/ostorlab/agent2"
+    assert agents["pageInfo"]["count"] == 3
+    assert agents["pageInfo"]["numPages"] == 2
+    assert agents["pageInfo"]["hasNext"] is True
+    assert agents["pageInfo"]["hasPrevious"] is False