diff --git a/examples/agent_pool.py b/examples/agent_pool.py index bbaf14e..432c706 100644 --- a/examples/agent_pool.py +++ b/examples/agent_pool.py @@ -3,13 +3,14 @@ This example demonstrates: 1. Agent Pool CRUD operations (Create, Read, Update, Delete) 2. Agent token creation and management -3. Using the organization SDK client +3. Workspace assignment using assign_to_workspaces and remove_from_workspaces 4. Proper error handling Make sure to set the following environment variables: - TFE_TOKEN: Your Terraform Cloud/Enterprise API token - TFE_ADDRESS: Your Terraform Cloud/Enterprise URL (optional, defaults to https://app.terraform.io) - TFE_ORG: Your organization name +- TFE_WORKSPACE_ID: A workspace ID for testing workspace assignment (optional) Usage: export TFE_TOKEN="your-token-here" @@ -24,8 +25,10 @@ from pytfe.errors import NotFound from pytfe.models import ( AgentPoolAllowedWorkspacePolicy, + AgentPoolAssignToWorkspacesOptions, AgentPoolCreateOptions, AgentPoolListOptions, + AgentPoolRemoveFromWorkspacesOptions, AgentPoolUpdateOptions, AgentTokenCreateOptions, ) @@ -37,6 +40,9 @@ def main(): token = os.environ.get("TFE_TOKEN") org = os.environ.get("TFE_ORG") address = os.environ.get("TFE_ADDRESS", "https://app.terraform.io") + workspace_id = os.environ.get( + "TFE_WORKSPACE_ID" + ) # optional, for workspace assignment if not token: print("TFE_TOKEN environment variable is required") @@ -96,7 +102,27 @@ def main(): updated_pool = client.agent_pools.update(new_pool.id, update_options) print(f"Updated agent pool name to: {updated_pool.name}") - # Example 5: Create an agent token + # Example 5: Workspace assignment + # assign_to_workspaces sends PATCH /agent-pools/:id with relationships.allowed-workspaces + # remove_from_workspaces sends PATCH /agent-pools/:id with relationships.excluded-workspaces + if workspace_id: + print("\n Assigning workspace to agent pool...") + updated_pool = client.agent_pools.assign_to_workspaces( + new_pool.id, + AgentPoolAssignToWorkspacesOptions(workspace_ids=[workspace_id]), + ) + print(f" Assigned workspace {workspace_id} to pool {updated_pool.name}") + + print("\n Removing workspace from agent pool...") + updated_pool = client.agent_pools.remove_from_workspaces( + new_pool.id, + AgentPoolRemoveFromWorkspacesOptions(workspace_ids=[workspace_id]), + ) + print(f" Removed workspace {workspace_id} from pool {updated_pool.name}") + else: + print("\n Skipping workspace assignment (set TFE_WORKSPACE_ID to test)") + + # Example 6: Create an agent token print("\n Creating agent token...") token_options = AgentTokenCreateOptions( description="SDK example token" # Optional description @@ -107,7 +133,7 @@ def main(): if agent_token.token: print(f" Token (first 10 chars): {agent_token.token[:10]}...") - # Example 6: List agent tokens + # Example 7: List agent tokens print("\n Listing agent tokens...") tokens = client.agent_tokens.list(new_pool.id) @@ -117,7 +143,7 @@ def main(): for token in token_list: print(f" - {token.description or 'No description'} (ID: {token.id})") - # Example 7: Clean up - delete the token and pool + # Example 8: Clean up - delete the token and pool print("\n Cleaning up...") client.agent_tokens.delete(agent_token.id) print("Deleted agent token") diff --git a/src/pytfe/models/agent.py b/src/pytfe/models/agent.py index 48c3055..08f40e1 100644 --- a/src/pytfe/models/agent.py +++ b/src/pytfe/models/agent.py @@ -82,6 +82,10 @@ class AgentPoolCreateOptions(BaseModel): organization_scoped: bool | None = None # Optional: Allowed workspace policy allowed_workspace_policy: AgentPoolAllowedWorkspacePolicy | None = None + # Optional: IDs of workspaces allowed to use this pool (sent as relationships.allowed-workspaces) + allowed_workspace_ids: list[str] = Field(default_factory=list) + # Optional: IDs of workspaces excluded from this pool (sent as relationships.excluded-workspaces) + excluded_workspace_ids: list[str] = Field(default_factory=list) class AgentPoolUpdateOptions(BaseModel): @@ -93,6 +97,10 @@ class AgentPoolUpdateOptions(BaseModel): organization_scoped: bool | None = None # Optional: Allowed workspace policy allowed_workspace_policy: AgentPoolAllowedWorkspacePolicy | None = None + # Optional: Full replacement list of workspace IDs allowed to use this pool + allowed_workspace_ids: list[str] = Field(default_factory=list) + # Optional: Full replacement list of workspace IDs excluded from this pool + excluded_workspace_ids: list[str] = Field(default_factory=list) class AgentPoolReadOptions(BaseModel): diff --git a/src/pytfe/resources/agent_pools.py b/src/pytfe/resources/agent_pools.py index e0ff776..7a90f2c 100644 --- a/src/pytfe/resources/agent_pools.py +++ b/src/pytfe/resources/agent_pools.py @@ -203,7 +203,27 @@ def create(self, organization: str, options: AgentPoolCreateOptions) -> AgentPoo options.allowed_workspace_policy.value ) - payload = {"data": {"type": "agent-pools", "attributes": attributes}} + relationships: dict[str, Any] = {} + if options.allowed_workspace_ids: + relationships["allowed-workspaces"] = { + "data": [ + {"type": "workspaces", "id": ws_id} + for ws_id in options.allowed_workspace_ids + ] + } + if options.excluded_workspace_ids: + relationships["excluded-workspaces"] = { + "data": [ + {"type": "workspaces", "id": ws_id} + for ws_id in options.excluded_workspace_ids + ] + } + + payload: dict[str, Any] = { + "data": {"type": "agent-pools", "attributes": attributes} + } + if relationships: + payload["data"]["relationships"] = relationships response = self.t.request("POST", path, json_body=payload) data = response.json()["data"] @@ -320,13 +340,31 @@ def update(self, agent_pool_id: str, options: AgentPoolUpdateOptions) -> AgentPo options.allowed_workspace_policy.value ) - payload = { + relationships: dict[str, Any] = {} + if options.allowed_workspace_ids: + relationships["allowed-workspaces"] = { + "data": [ + {"type": "workspaces", "id": ws_id} + for ws_id in options.allowed_workspace_ids + ] + } + if options.excluded_workspace_ids: + relationships["excluded-workspaces"] = { + "data": [ + {"type": "workspaces", "id": ws_id} + for ws_id in options.excluded_workspace_ids + ] + } + + payload: dict[str, Any] = { "data": { "type": "agent-pools", "id": agent_pool_id, "attributes": attributes, } } + if relationships: + payload["data"]["relationships"] = relationships response = self.t.request("PATCH", path, json_body=payload) data = response.json()["data"] @@ -371,13 +409,20 @@ def delete(self, agent_pool_id: str) -> None: def assign_to_workspaces( self, agent_pool_id: str, options: AgentPoolAssignToWorkspacesOptions - ) -> None: - """Assign an agent pool to workspaces. + ) -> AgentPool: + """Assign an agent pool to workspaces by updating the allowed-workspaces + relationship via PATCH /agent-pools/:id. + + The provided workspace IDs become the new complete list of allowed + workspaces for this pool (full replacement, not append). Args: 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 @@ -388,26 +433,67 @@ def assign_to_workspaces( if not options.workspace_ids: raise ValueError("At least one workspace ID is required") - path = f"/api/v2/agent-pools/{agent_pool_id}/relationships/workspaces" - - # Create data payload with workspace references - workspace_data = [] for workspace_id in options.workspace_ids: if not valid_string_id(workspace_id): raise ValueError(f"Invalid workspace ID: {workspace_id}") - workspace_data.append({"type": "workspaces", "id": workspace_id}) - payload = {"data": workspace_data} - self.t.request("POST", path, json_body=payload) + path = f"/api/v2/agent-pools/{agent_pool_id}" + payload: dict[str, Any] = { + "data": { + "type": "agent-pools", + "id": agent_pool_id, + "attributes": {}, + "relationships": { + "allowed-workspaces": { + "data": [ + {"type": "workspaces", "id": ws_id} + for ws_id in options.workspace_ids + ] + } + }, + } + } + 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: - """Remove an agent pool from workspaces. + ) -> AgentPool: + """Exclude workspaces from an agent pool by updating the excluded-workspaces + relationship via PATCH /agent-pools/:id. + + Use this for organization-scoped pools where most workspaces are allowed + but you want to block specific ones. The provided list becomes the new + complete excluded-workspaces list (full replacement, not append). Args: agent_pool_id: Agent pool ID - options: Removal options containing workspace IDs + options: Removal options containing workspace IDs to exclude + + Returns: + Updated AgentPool object Raises: ValueError: If parameters are invalid @@ -419,14 +505,47 @@ def remove_from_workspaces( if not options.workspace_ids: raise ValueError("At least one workspace ID is required") - path = f"/api/v2/agent-pools/{agent_pool_id}/relationships/workspaces" - - # Create data payload with workspace references - workspace_data = [] for workspace_id in options.workspace_ids: if not valid_string_id(workspace_id): raise ValueError(f"Invalid workspace ID: {workspace_id}") - workspace_data.append({"type": "workspaces", "id": workspace_id}) - payload = {"data": workspace_data} - self.t.request("DELETE", path, json_body=payload) + path = f"/api/v2/agent-pools/{agent_pool_id}" + payload: dict[str, Any] = { + "data": { + "type": "agent-pools", + "id": agent_pool_id, + "attributes": {}, + "relationships": { + "excluded-workspaces": { + "data": [ + {"type": "workspaces", "id": ws_id} + for ws_id in options.workspace_ids + ] + } + }, + } + } + 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"]), + ) diff --git a/tests/units/test_agent_pools.py b/tests/units/test_agent_pools.py index a113ac6..5022ffd 100644 --- a/tests/units/test_agent_pools.py +++ b/tests/units/test_agent_pools.py @@ -6,6 +6,7 @@ 3. Agent token management 4. Request building and parameter handling 5. Response parsing and error handling +6. Workspace assignment (assign_to_workspaces / remove_from_workspaces bug fix) Run with: pytest tests/units/test_agent_pools.py -v @@ -19,8 +20,10 @@ from pytfe.models.agent import ( AgentPool, AgentPoolAllowedWorkspacePolicy, + AgentPoolAssignToWorkspacesOptions, AgentPoolCreateOptions, AgentPoolListOptions, + AgentPoolRemoveFromWorkspacesOptions, AgentPoolUpdateOptions, AgentTokenCreateOptions, ) @@ -85,6 +88,26 @@ def test_agent_pool_create_options(self): == AgentPoolAllowedWorkspacePolicy.SPECIFIC_WORKSPACES ) + def test_agent_pool_create_options_workspace_ids(self): + """Test AgentPoolCreateOptions with allowed/excluded workspace IDs (bug fix)""" + options = AgentPoolCreateOptions( + name="scoped-pool", + organization_scoped=False, + allowed_workspace_ids=["ws-aaa", "ws-bbb"], + excluded_workspace_ids=["ws-ccc"], + ) + assert options.allowed_workspace_ids == ["ws-aaa", "ws-bbb"] + assert options.excluded_workspace_ids == ["ws-ccc"] + + def test_agent_pool_update_options_workspace_ids(self): + """Test AgentPoolUpdateOptions with allowed/excluded workspace IDs (bug fix)""" + options = AgentPoolUpdateOptions( + allowed_workspace_ids=["ws-aaa"], + excluded_workspace_ids=["ws-bbb"], + ) + assert options.allowed_workspace_ids == ["ws-aaa"] + assert options.excluded_workspace_ids == ["ws-bbb"] + class TestAgentPoolOperations: """Test agent pool CRUD operations""" @@ -121,7 +144,6 @@ def test_list_agent_pools(self, agent_pools_service, mock_transport): } mock_transport.request.return_value.json.return_value = mock_response - agent_pools = list(agent_pools_service.list("test-org")) assert len(agent_pools) == 1 @@ -171,7 +193,6 @@ def test_create_agent_pool(self, agent_pools_service, mock_transport): } mock_transport.request.return_value.json.return_value = mock_response - options = AgentPoolCreateOptions( name="new-pool", organization_scoped=True, @@ -207,7 +228,6 @@ def test_read_agent_pool(self, agent_pools_service, mock_transport): } mock_transport.request.return_value.json.return_value = mock_response - agent_pool = agent_pools_service.read("apool-123456789abcdef0") assert agent_pool.id == "apool-123456789abcdef0" @@ -238,9 +258,7 @@ def test_update_agent_pool(self, agent_pools_service, mock_transport): } mock_transport.request.return_value.json.return_value = mock_response - options = AgentPoolUpdateOptions(name="updated-pool", organization_scoped=False) - agent_pool = agent_pools_service.update("apool-123456789abcdef0", options) assert agent_pool.id == "apool-123456789abcdef0" @@ -263,6 +281,96 @@ def test_delete_agent_pool(self, agent_pools_service, mock_transport): assert call_args[0][0] == "DELETE" assert "agent-pools/apool-123456789abcdef0" in call_args[0][1] + def test_assign_to_workspaces(self, agent_pools_service, mock_transport): + """assign_to_workspaces must PATCH /agent-pools/:id with relationships.allowed-workspaces. + + Previously (broken): POST /agent-pools/:id/relationships/workspaces -> 404 + Fixed: PATCH /agent-pools/:id with relationships.allowed-workspaces body + """ + pool_id = "apool-123456789abcdef0" + ws_id = "ws-aaaaaaaaaaaaaaa1" + + 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" + # Must target the pool URL, not a /relationships/workspaces sub-resource + assert call_args[0][1] == f"/api/v2/agent-pools/{pool_id}" + # Payload must use relationships.allowed-workspaces + body = call_args[1]["json_body"]["data"] + assert body["type"] == "agent-pools" + assert body["id"] == pool_id + ws_data = body["relationships"]["allowed-workspaces"]["data"] + assert ws_data[0]["id"] == ws_id + assert ws_data[0]["type"] == "workspaces" + + def test_remove_from_workspaces(self, agent_pools_service, mock_transport): + """remove_from_workspaces must PATCH /agent-pools/:id with relationships.excluded-workspaces. + + Previously (broken): DELETE /agent-pools/:id/relationships/workspaces -> 404 + Fixed: PATCH /agent-pools/:id with relationships.excluded-workspaces body + """ + pool_id = "apool-123456789abcdef0" + ws_id = "ws-aaaaaaaaaaaaaaa1" + + 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" + # Must target the pool URL, not a /relationships/workspaces sub-resource + assert call_args[0][1] == f"/api/v2/agent-pools/{pool_id}" + # Payload must use relationships.excluded-workspaces + body = call_args[1]["json_body"]["data"] + assert body["type"] == "agent-pools" + assert body["id"] == pool_id + ws_data = body["relationships"]["excluded-workspaces"]["data"] + assert ws_data[0]["id"] == ws_id + assert ws_data[0]["type"] == "workspaces" + class TestAgentTokenOperations: """Test agent token operations""" @@ -297,7 +405,6 @@ def test_list_agent_tokens(self, agent_tokens_service, mock_transport): } mock_transport.request.return_value.json.return_value = mock_response - tokens = list(agent_tokens_service.list("apool-123456789abcdef0")) assert len(tokens) == 1 @@ -328,7 +435,6 @@ def test_create_agent_token(self, agent_tokens_service, mock_transport): } mock_transport.request.return_value.json.return_value = mock_response - options = AgentTokenCreateOptions(description="New token") token = agent_tokens_service.create("apool-123456789abcdef0", options) @@ -361,7 +467,6 @@ def test_read_agent_token(self, agent_tokens_service, mock_transport): } mock_transport.request.return_value.json.return_value = mock_response - token = agent_tokens_service.read("at-123456789abcdef0") assert token.id == "at-123456789abcdef0"