agent_pools: return AgentPool from assign/remove workspace operations; update tests and example

Author: KshitijaChoudhari

examples/agent_pool.py

@@ -107,18 +107,18 @@ def main():
         # remove_from_workspaces sends PATCH /agent-pools/:id with relationships.excluded-workspaces
         if workspace_id:
             print("\n Assigning workspace to agent pool...")
-            client.agent_pools.assign_to_workspaces(
+            updated_pool = client.agent_pools.assign_to_workspaces(
                 new_pool.id,
                 AgentPoolAssignToWorkspacesOptions(workspace_ids=[workspace_id]),
             )
-            print(f"  Assigned workspace {workspace_id} to pool")
+            print(f"  Assigned workspace {workspace_id} to pool {updated_pool.name}")
 
             print("\n Removing workspace from agent pool...")
-            client.agent_pools.remove_from_workspaces(
+            updated_pool = client.agent_pools.remove_from_workspaces(
                 new_pool.id,
                 AgentPoolRemoveFromWorkspacesOptions(workspace_ids=[workspace_id]),
             )
-            print(f"  Excluded workspace {workspace_id} from pool")
+            print(f"  Removed workspace {workspace_id} from pool {updated_pool.name}")
         else:
             print("\n Skipping workspace assignment (set TFE_WORKSPACE_ID to test)")
 

src/pytfe/resources/agent_pools.py

@@ -409,7 +409,7 @@ def delete(self, agent_pool_id: str) -> None:
 
     def assign_to_workspaces(
         self, agent_pool_id: str, options: AgentPoolAssignToWorkspacesOptions
-    ) -> None:
+    ) -> AgentPool:
         """Assign an agent pool to workspaces by updating the allowed-workspaces
         relationship via PATCH /agent-pools/:id.
 
@@ -420,6 +420,9 @@ def assign_to_workspaces(
             agent_pool_id: Agent pool ID
             options: Assignment options containing workspace IDs
 
+        Returns:
+            Updated AgentPool object
+
         Raises:
             ValueError: If parameters are invalid
             TFEError: If API request fails
@@ -450,11 +453,34 @@ def assign_to_workspaces(
                 },
             }
         }
-        self.t.request("PATCH", path, json_body=payload)
+        response = self.t.request("PATCH", path, json_body=payload)
+        data = response.json()["data"]
+
+        # Extract agent pool data from response
+        attr = data.get("attributes", {}) or {}
+        agent_pool_data = {
+            "id": _safe_str(data.get("id")),
+            "name": _safe_str(attr.get("name")),
+            "created_at": attr.get("created-at"),
+            "organization_scoped": attr.get("organization-scoped"),
+            "allowed_workspace_policy": attr.get("allowed-workspace-policy"),
+            "agent_count": attr.get("agent-count", 0),
+        }
+
+        return AgentPool(
+            id=_safe_str(agent_pool_data["id"]) or "",
+            name=_safe_str(agent_pool_data["name"]),
+            created_at=cast(Any, agent_pool_data["created_at"]),
+            organization_scoped=_safe_bool(agent_pool_data["organization_scoped"]),
+            allowed_workspace_policy=_safe_workspace_policy(
+                agent_pool_data["allowed_workspace_policy"]
+            ),
+            agent_count=_safe_int(agent_pool_data["agent_count"]),
+        )
 
     def remove_from_workspaces(
         self, agent_pool_id: str, options: AgentPoolRemoveFromWorkspacesOptions
-    ) -> None:
+    ) -> AgentPool:
         """Exclude workspaces from an agent pool by updating the excluded-workspaces
         relationship via PATCH /agent-pools/:id.
 
@@ -466,6 +492,9 @@ def remove_from_workspaces(
             agent_pool_id: Agent pool ID
             options: Removal options containing workspace IDs to exclude
 
+        Returns:
+            Updated AgentPool object
+
         Raises:
             ValueError: If parameters are invalid
             TFEError: If API request fails
@@ -496,4 +525,27 @@ def remove_from_workspaces(
                 },
             }
         }
-        self.t.request("PATCH", path, json_body=payload)
+        response = self.t.request("PATCH", path, json_body=payload)
+        data = response.json()["data"]
+
+        # Extract agent pool data from response
+        attr = data.get("attributes", {}) or {}
+        agent_pool_data = {
+            "id": _safe_str(data.get("id")),
+            "name": _safe_str(attr.get("name")),
+            "created_at": attr.get("created-at"),
+            "organization_scoped": attr.get("organization-scoped"),
+            "allowed_workspace_policy": attr.get("allowed-workspace-policy"),
+            "agent_count": attr.get("agent-count", 0),
+        }
+
+        return AgentPool(
+            id=_safe_str(agent_pool_data["id"]) or "",
+            name=_safe_str(agent_pool_data["name"]),
+            created_at=cast(Any, agent_pool_data["created_at"]),
+            organization_scoped=_safe_bool(agent_pool_data["organization_scoped"]),
+            allowed_workspace_policy=_safe_workspace_policy(
+                agent_pool_data["allowed_workspace_policy"]
+            ),
+            agent_count=_safe_int(agent_pool_data["agent_count"]),
+        )

tests/units/test_agent_pools.py

@@ -290,11 +290,29 @@ def test_assign_to_workspaces(self, agent_pools_service, mock_transport):
         pool_id = "apool-123456789abcdef0"
         ws_id = "ws-aaaaaaaaaaaaaaa1"
 
-        agent_pools_service.assign_to_workspaces(
+        mock_response = {
+            "data": {
+                "id": pool_id,
+                "type": "agent-pools",
+                "attributes": {
+                    "name": "test-pool",
+                    "created-at": "2023-01-01T00:00:00Z",
+                    "organization-scoped": True,
+                    "allowed-workspace-policy": "all-workspaces",
+                    "agent-count": 0,
+                },
+            }
+        }
+        mock_transport.request.return_value.json.return_value = mock_response
+
+        agent_pool = agent_pools_service.assign_to_workspaces(
             pool_id,
             AgentPoolAssignToWorkspacesOptions(workspace_ids=[ws_id]),
         )
 
+        assert agent_pool.id == pool_id
+        assert agent_pool.name == "test-pool"
+
         call_args = mock_transport.request.call_args
         # Must be PATCH, not POST
         assert call_args[0][0] == "PATCH"
@@ -317,11 +335,29 @@ def test_remove_from_workspaces(self, agent_pools_service, mock_transport):
         pool_id = "apool-123456789abcdef0"
         ws_id = "ws-aaaaaaaaaaaaaaa1"
 
-        agent_pools_service.remove_from_workspaces(
+        mock_response = {
+            "data": {
+                "id": pool_id,
+                "type": "agent-pools",
+                "attributes": {
+                    "name": "test-pool",
+                    "created-at": "2023-01-01T00:00:00Z",
+                    "organization-scoped": True,
+                    "allowed-workspace-policy": "all-workspaces",
+                    "agent-count": 0,
+                },
+            }
+        }
+        mock_transport.request.return_value.json.return_value = mock_response
+
+        agent_pool = agent_pools_service.remove_from_workspaces(
             pool_id,
             AgentPoolRemoveFromWorkspacesOptions(workspace_ids=[ws_id]),
         )
 
+        assert agent_pool.id == pool_id
+        assert agent_pool.name == "test-pool"
+
         call_args = mock_transport.request.call_args
         # Must be PATCH, not DELETE
         assert call_args[0][0] == "PATCH"