From 67206928540917433ea88a98319529779bb2d81e Mon Sep 17 00:00:00 2001 From: Tanya Singh Date: Thu, 21 May 2026 12:45:45 +0530 Subject: [PATCH 1/2] Add agent pool relationship update APIs --- examples/agent_pool.py | 63 +++++- src/pytfe/models/__init__.py | 6 + src/pytfe/models/agent.py | 40 ++++ src/pytfe/resources/agent_pools.py | 121 ++++++++++- tests/units/test_agent_pools.py | 311 +++++++++++++++++++++++++++++ 5 files changed, 539 insertions(+), 2 deletions(-) diff --git a/examples/agent_pool.py b/examples/agent_pool.py index bcb04ae3..1d5c37c6 100644 --- a/examples/agent_pool.py +++ b/examples/agent_pool.py @@ -7,13 +7,17 @@ 1. Agent Pool CRUD operations (Create, Read, Update, Delete) 2. Agent token creation and management 3. Workspace assignment using assign_to_workspaces and remove_from_workspaces -4. Proper error handling +4. Project assignment using update_allowed_projects (Go SDK parity) +5. Dedicated relationship update methods: update_allowed_workspaces, + update_allowed_projects, update_excluded_workspaces +6. 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) +- TFE_PROJECT_ID: A project ID for testing project assignment (optional) Usage: export TFE_TOKEN="your-token-here" @@ -27,9 +31,12 @@ from pytfe import TFEClient, TFEConfig from pytfe.errors import NotFound from pytfe.models import ( + AgentPoolAllowedProjectsUpdateOptions, AgentPoolAllowedWorkspacePolicy, + AgentPoolAllowedWorkspacesUpdateOptions, AgentPoolAssignToWorkspacesOptions, AgentPoolCreateOptions, + AgentPoolExcludedWorkspacesUpdateOptions, AgentPoolListOptions, AgentPoolRemoveFromWorkspacesOptions, AgentPoolUpdateOptions, @@ -46,6 +53,7 @@ def main(): workspace_id = os.environ.get( "TFE_WORKSPACE_ID" ) # optional, for workspace assignment + project_id = os.environ.get("TFE_PROJECT_ID") # optional, for project assignment if not token: print("TFE_TOKEN environment variable is required") @@ -122,9 +130,62 @@ def main(): AgentPoolRemoveFromWorkspacesOptions(workspace_ids=[workspace_id]), ) print(f" Removed workspace {workspace_id} from pool {updated_pool.name}") + + # Example 5b: Dedicated update methods + # update_allowed_workspaces / update_excluded_workspaces send the + # relationship array unconditionally — even an empty list — so they can + # also CLEAR the relationship. + print( + "\n Using update_allowed_workspaces (dedicated, supports clearing)..." + ) + updated_pool = client.agent_pools.update_allowed_workspaces( + new_pool.id, + AgentPoolAllowedWorkspacesUpdateOptions(workspace_ids=[workspace_id]), + ) + print( + f" Set allowed-workspaces to [{workspace_id}] on pool {updated_pool.name}" + ) + + print("\n Clearing allowed-workspaces via update_allowed_workspaces...") + updated_pool = client.agent_pools.update_allowed_workspaces( + new_pool.id, + AgentPoolAllowedWorkspacesUpdateOptions(workspace_ids=[]), + ) + print(f" Cleared allowed-workspaces on pool {updated_pool.name}") + + print( + "\n Using update_excluded_workspaces (dedicated, supports clearing)..." + ) + updated_pool = client.agent_pools.update_excluded_workspaces( + new_pool.id, + AgentPoolExcludedWorkspacesUpdateOptions(workspace_ids=[workspace_id]), + ) + print( + f" Set excluded-workspaces to [{workspace_id}] on pool {updated_pool.name}" + ) else: print("\n Skipping workspace assignment (set TFE_WORKSPACE_ID to test)") + # Example 5c: Project assignment (parity — AllowedProjects relationship) + if project_id: + print("\n Assigning project to agent pool (update_allowed_projects)...") + updated_pool = client.agent_pools.update_allowed_projects( + new_pool.id, + AgentPoolAllowedProjectsUpdateOptions(project_ids=[project_id]), + ) + print( + f" Set allowed-projects to [{project_id}] on pool {updated_pool.name}" + ) + + print("\n Clearing allowed-projects via update_allowed_projects...") + updated_pool = client.agent_pools.update_allowed_projects( + new_pool.id, + AgentPoolAllowedProjectsUpdateOptions(project_ids=[]), + ) + print(f" Cleared allowed-projects on pool {updated_pool.name}") + else: + print("\n Skipping project assignment (set TFE_PROJECT_ID to test)") + # Example 6: Create an agent token print("\n Creating agent token...") token_options = AgentTokenCreateOptions( diff --git a/src/pytfe/models/__init__.py b/src/pytfe/models/__init__.py index fd263c1d..5f38f3e5 100644 --- a/src/pytfe/models/__init__.py +++ b/src/pytfe/models/__init__.py @@ -8,9 +8,12 @@ Agent, AgentListOptions, AgentPool, + AgentPoolAllowedProjectsUpdateOptions, AgentPoolAllowedWorkspacePolicy, + AgentPoolAllowedWorkspacesUpdateOptions, AgentPoolAssignToWorkspacesOptions, AgentPoolCreateOptions, + AgentPoolExcludedWorkspacesUpdateOptions, AgentPoolListOptions, AgentPoolReadOptions, AgentPoolRemoveFromWorkspacesOptions, @@ -463,9 +466,12 @@ # Agent & pools "Agent", "AgentPool", + "AgentPoolAllowedProjectsUpdateOptions", "AgentPoolAllowedWorkspacePolicy", + "AgentPoolAllowedWorkspacesUpdateOptions", "AgentPoolAssignToWorkspacesOptions", "AgentPoolCreateOptions", + "AgentPoolExcludedWorkspacesUpdateOptions", "AgentPoolListOptions", "AgentPoolReadOptions", "AgentPoolRemoveFromWorkspacesOptions", diff --git a/src/pytfe/models/agent.py b/src/pytfe/models/agent.py index d0751ea1..c9761203 100644 --- a/src/pytfe/models/agent.py +++ b/src/pytfe/models/agent.py @@ -74,6 +74,10 @@ class AgentPoolListOptions(BaseModel): include: list[str] | None = None # Optional: Filter by allowed workspace policy allowed_workspace_policy: AgentPoolAllowedWorkspacePolicy | None = None + # Optional: String (workspace name) used to filter the results + allowed_workspaces_name: str | None = None + # Optional: String (project name) used to filter the results + allowed_projects_name: str | None = None class AgentPoolCreateOptions(BaseModel): @@ -87,6 +91,8 @@ class AgentPoolCreateOptions(BaseModel): 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 projects allowed to use this pool (sent as relationships.allowed-projects) + allowed_project_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) @@ -102,6 +108,8 @@ class AgentPoolUpdateOptions(BaseModel): 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 project IDs allowed to use this pool + allowed_project_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) @@ -128,6 +136,38 @@ class AgentPoolRemoveFromWorkspacesOptions(BaseModel): workspace_ids: list[str] = Field(default_factory=list) +# Dedicated relationship-update options. +# Unlike the main create/update options, these always send their relationship array in the +# payload — even when empty — so that the caller can clear existing relationships. + + +class AgentPoolAllowedWorkspacesUpdateOptions(BaseModel): + """Options for updating the allowed-workspaces relationship on an agent pool. + + Supports full replacement including clearing (pass an empty list). + """ + + workspace_ids: list[str] = Field(default_factory=list) + + +class AgentPoolAllowedProjectsUpdateOptions(BaseModel): + """Options for updating the allowed-projects relationship on an agent pool. + + Supports full replacement including clearing (pass an empty list). + """ + + project_ids: list[str] = Field(default_factory=list) + + +class AgentPoolExcludedWorkspacesUpdateOptions(BaseModel): + """Options for updating the excluded-workspaces relationship on an agent pool. + + Supports full replacement including clearing (pass an empty list). + """ + + workspace_ids: list[str] = Field(default_factory=list) + + # Agent Options diff --git a/src/pytfe/resources/agent_pools.py b/src/pytfe/resources/agent_pools.py index 47ffc8bf..344d3628 100644 --- a/src/pytfe/resources/agent_pools.py +++ b/src/pytfe/resources/agent_pools.py @@ -9,14 +9,17 @@ from __future__ import annotations -from collections.abc import Iterator +from collections.abc import Iterator, Sequence from typing import Any, cast from ..models.agent import ( AgentPool, + AgentPoolAllowedProjectsUpdateOptions, AgentPoolAllowedWorkspacePolicy, + AgentPoolAllowedWorkspacesUpdateOptions, AgentPoolAssignToWorkspacesOptions, AgentPoolCreateOptions, + AgentPoolExcludedWorkspacesUpdateOptions, AgentPoolListOptions, AgentPoolReadOptions, AgentPoolRemoveFromWorkspacesOptions, @@ -143,6 +146,12 @@ def list( params["filter[allowed_workspace_policy]"] = ( options.allowed_workspace_policy.value ) + if options.allowed_workspaces_name: + params["filter[allowed_workspaces][name]"] = ( + options.allowed_workspaces_name + ) + if options.allowed_projects_name: + params["filter[allowed_projects][name]"] = options.allowed_projects_name items_iter = self._list(path, params=params) @@ -214,6 +223,13 @@ def create(self, organization: str, options: AgentPoolCreateOptions) -> AgentPoo for ws_id in options.allowed_workspace_ids ] } + if options.allowed_project_ids: + relationships["allowed-projects"] = { + "data": [ + {"type": "projects", "id": proj_id} + for proj_id in options.allowed_project_ids + ] + } if options.excluded_workspace_ids: relationships["excluded-workspaces"] = { "data": [ @@ -351,6 +367,13 @@ def update(self, agent_pool_id: str, options: AgentPoolUpdateOptions) -> AgentPo for ws_id in options.allowed_workspace_ids ] } + if options.allowed_project_ids: + relationships["allowed-projects"] = { + "data": [ + {"type": "projects", "id": proj_id} + for proj_id in options.allowed_project_ids + ] + } if options.excluded_workspace_ids: relationships["excluded-workspaces"] = { "data": [ @@ -552,3 +575,99 @@ def remove_from_workspaces( ), agent_count=_safe_int(agent_pool_data["agent_count"]), ) + + def _patch_relationship( + self, + agent_pool_id: str, + relationship_key: str, + items_data: Sequence[dict[str, Any]], + ) -> AgentPool: + """Internal helper: PATCH agent-pools/:id with a single relationship key. + + Always sends the relationship array in the payload even when empty, which + allows callers to clear an existing relationship by passing an empty list. + This mirrors the Go SDK updateArrayAttribute helper. + """ + if not valid_string_id(agent_pool_id): + raise ValueError("Agent pool ID is required and must be valid") + + path = f"/api/v2/agent-pools/{agent_pool_id}" + payload: dict[str, Any] = { + "data": { + "type": "agent-pools", + "id": agent_pool_id, + "attributes": {}, + "relationships": {relationship_key: {"data": items_data}}, + } + } + response = self.t.request("PATCH", path, json_body=payload) + data = response.json()["data"] + + attr = data.get("attributes", {}) or {} + return AgentPool( + id=_safe_str(data.get("id")) or "", + name=_safe_str(attr.get("name")), + created_at=cast(Any, attr.get("created-at")), + organization_scoped=_safe_bool(attr.get("organization-scoped")), + allowed_workspace_policy=_safe_workspace_policy( + attr.get("allowed-workspace-policy") + ), + agent_count=_safe_int(attr.get("agent-count", 0)), + ) + + def update_allowed_workspaces( + self, agent_pool_id: str, options: AgentPoolAllowedWorkspacesUpdateOptions + ) -> AgentPool: + """Update the complete allowed-workspaces list for an agent pool. + + Args: + agent_pool_id: Agent pool ID + options: Options containing the new full list of workspace IDs + + Returns: + Updated AgentPool object + + Raises: + ValueError: If agent_pool_id is invalid + TFEError: If API request fails + """ + items = [{"type": "workspaces", "id": ws_id} for ws_id in options.workspace_ids] + return self._patch_relationship(agent_pool_id, "allowed-workspaces", items) + + def update_allowed_projects( + self, agent_pool_id: str, options: AgentPoolAllowedProjectsUpdateOptions + ) -> AgentPool: + """Update the complete allowed-projects list for an agent pool. + + Args: + agent_pool_id: Agent pool ID + options: Options containing the new full list of project IDs + + Returns: + Updated AgentPool object + + Raises: + ValueError: If agent_pool_id is invalid + TFEError: If API request fails + """ + items = [{"type": "projects", "id": proj_id} for proj_id in options.project_ids] + return self._patch_relationship(agent_pool_id, "allowed-projects", items) + + def update_excluded_workspaces( + self, agent_pool_id: str, options: AgentPoolExcludedWorkspacesUpdateOptions + ) -> AgentPool: + """Update the complete excluded-workspaces list for an agent pool. + + Args: + agent_pool_id: Agent pool ID + options: Options containing the new full list of workspace IDs to exclude + + Returns: + Updated AgentPool object + + Raises: + ValueError: If agent_pool_id is invalid + TFEError: If API request fails + """ + items = [{"type": "workspaces", "id": ws_id} for ws_id in options.workspace_ids] + return self._patch_relationship(agent_pool_id, "excluded-workspaces", items) diff --git a/tests/units/test_agent_pools.py b/tests/units/test_agent_pools.py index f797f612..e0bfed65 100644 --- a/tests/units/test_agent_pools.py +++ b/tests/units/test_agent_pools.py @@ -22,9 +22,12 @@ from pytfe.errors import AuthError, NotFound, ValidationError from pytfe.models.agent import ( AgentPool, + AgentPoolAllowedProjectsUpdateOptions, AgentPoolAllowedWorkspacePolicy, + AgentPoolAllowedWorkspacesUpdateOptions, AgentPoolAssignToWorkspacesOptions, AgentPoolCreateOptions, + AgentPoolExcludedWorkspacesUpdateOptions, AgentPoolListOptions, AgentPoolRemoveFromWorkspacesOptions, AgentPoolUpdateOptions, @@ -102,6 +105,15 @@ def test_agent_pool_create_options_workspace_ids(self): assert options.allowed_workspace_ids == ["ws-aaa", "ws-bbb"] assert options.excluded_workspace_ids == ["ws-ccc"] + def test_agent_pool_create_options_project_ids(self): + """Test AgentPoolCreateOptions with allowed_project_ids""" + options = AgentPoolCreateOptions( + name="scoped-pool", + organization_scoped=False, + allowed_project_ids=["prj-aaa", "prj-bbb"], + ) + assert options.allowed_project_ids == ["prj-aaa", "prj-bbb"] + def test_agent_pool_update_options_workspace_ids(self): """Test AgentPoolUpdateOptions with allowed/excluded workspace IDs (bug fix)""" options = AgentPoolUpdateOptions( @@ -111,6 +123,37 @@ def test_agent_pool_update_options_workspace_ids(self): assert options.allowed_workspace_ids == ["ws-aaa"] assert options.excluded_workspace_ids == ["ws-bbb"] + def test_agent_pool_update_options_project_ids(self): + """Test AgentPoolUpdateOptions with allowed_project_ids""" + options = AgentPoolUpdateOptions( + allowed_project_ids=["prj-aaa"], + ) + assert options.allowed_project_ids == ["prj-aaa"] + + def test_agent_pool_allowed_workspaces_update_options(self): + """Test AgentPoolAllowedWorkspacesUpdateOptions model""" + options = AgentPoolAllowedWorkspacesUpdateOptions( + workspace_ids=["ws-aaa", "ws-bbb"] + ) + assert options.workspace_ids == ["ws-aaa", "ws-bbb"] + + def test_agent_pool_allowed_projects_update_options(self): + """Test AgentPoolAllowedProjectsUpdateOptions model""" + options = AgentPoolAllowedProjectsUpdateOptions( + project_ids=["prj-aaa", "prj-bbb"] + ) + assert options.project_ids == ["prj-aaa", "prj-bbb"] + + def test_agent_pool_excluded_workspaces_update_options(self): + """Test AgentPoolExcludedWorkspacesUpdateOptions model""" + options = AgentPoolExcludedWorkspacesUpdateOptions(workspace_ids=["ws-aaa"]) + assert options.workspace_ids == ["ws-aaa"] + + def test_agent_pool_list_options_project_filter(self): + """Test AgentPoolListOptions now has allowed_projects_name filter""" + options = AgentPoolListOptions(allowed_projects_name="my-project") + assert options.allowed_projects_name == "my-project" + class TestAgentPoolOperations: """Test agent pool CRUD operations""" @@ -373,6 +416,274 @@ def test_remove_from_workspaces(self, agent_pools_service, mock_transport): assert ws_data[0]["id"] == ws_id assert ws_data[0]["type"] == "workspaces" + def test_create_agent_pool_with_project_ids( + self, agent_pools_service, mock_transport + ): + """Test creating an agent pool with allowed_project_ids serializes allowed-projects relationship""" + mock_response = { + "data": { + "id": "apool-123456789abcdef0", + "type": "agent-pools", + "attributes": { + "name": "scoped-pool", + "created-at": "2023-01-01T00:00:00Z", + "organization-scoped": False, + "allowed-workspace-policy": "specific-workspaces", + "agent-count": 0, + }, + } + } + mock_transport.request.return_value.json.return_value = mock_response + options = AgentPoolCreateOptions( + name="scoped-pool", + organization_scoped=False, + allowed_project_ids=["prj-aaa", "prj-bbb"], + ) + + agent_pools_service.create("test-org", options) + + call_args = mock_transport.request.call_args + assert call_args[0][0] == "POST" + body = call_args[1]["json_body"]["data"] + proj_data = body["relationships"]["allowed-projects"]["data"] + assert len(proj_data) == 2 + assert proj_data[0] == {"type": "projects", "id": "prj-aaa"} + assert proj_data[1] == {"type": "projects", "id": "prj-bbb"} + + def test_update_agent_pool_with_project_ids( + self, agent_pools_service, mock_transport + ): + """Test updating an agent pool with allowed_project_ids serializes allowed-projects relationship""" + pool_id = "apool-123456789abcdef0" + mock_response = { + "data": { + "id": pool_id, + "type": "agent-pools", + "attributes": { + "name": "test-pool", + "created-at": "2023-01-01T00:00:00Z", + "organization-scoped": False, + "allowed-workspace-policy": "specific-workspaces", + "agent-count": 0, + }, + } + } + mock_transport.request.return_value.json.return_value = mock_response + options = AgentPoolUpdateOptions(allowed_project_ids=["prj-aaa"]) + + agent_pools_service.update(pool_id, options) + + call_args = mock_transport.request.call_args + assert call_args[0][0] == "PATCH" + body = call_args[1]["json_body"]["data"] + proj_data = body["relationships"]["allowed-projects"]["data"] + assert proj_data == [{"type": "projects", "id": "prj-aaa"}] + + +class TestDedicatedRelationshipUpdateMethods: + """Test the Go SDK-parity dedicated update methods for relationships.""" + + @pytest.fixture + def mock_transport(self): + transport = Mock() + return transport + + @pytest.fixture + def agent_pools_service(self, mock_transport): + from pytfe.resources.agent_pools import AgentPools + + return AgentPools(mock_transport) + + def _pool_response(self, pool_id: str = "apool-123456789abcdef0") -> dict: + return { + "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, + }, + } + } + + def test_update_allowed_workspaces(self, agent_pools_service, mock_transport): + """update_allowed_workspaces PATCHes allowed-workspaces relationship.""" + pool_id = "apool-123456789abcdef0" + ws_id = "ws-aaaaaaaaaaaaaaa1" + mock_transport.request.return_value.json.return_value = self._pool_response( + pool_id + ) + + result = agent_pools_service.update_allowed_workspaces( + pool_id, + AgentPoolAllowedWorkspacesUpdateOptions(workspace_ids=[ws_id]), + ) + + assert result.id == pool_id + call_args = mock_transport.request.call_args + assert call_args[0][0] == "PATCH" + assert call_args[0][1] == f"/api/v2/agent-pools/{pool_id}" + 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 == [{"type": "workspaces", "id": ws_id}] + + def test_update_allowed_workspaces_empty_clears( + self, agent_pools_service, mock_transport + ): + """update_allowed_workspaces with empty list sends empty array to clear the relationship.""" + pool_id = "apool-123456789abcdef0" + mock_transport.request.return_value.json.return_value = self._pool_response( + pool_id + ) + + agent_pools_service.update_allowed_workspaces( + pool_id, + AgentPoolAllowedWorkspacesUpdateOptions(workspace_ids=[]), + ) + + call_args = mock_transport.request.call_args + body = call_args[1]["json_body"]["data"] + assert "allowed-workspaces" in body["relationships"] + assert body["relationships"]["allowed-workspaces"]["data"] == [] + + def test_update_allowed_projects(self, agent_pools_service, mock_transport): + """update_allowed_projects PATCHes allowed-projects relationship.""" + pool_id = "apool-123456789abcdef0" + proj_id = "prj-aaaaaaaaaaaaaaa1" + mock_transport.request.return_value.json.return_value = self._pool_response( + pool_id + ) + + result = agent_pools_service.update_allowed_projects( + pool_id, + AgentPoolAllowedProjectsUpdateOptions(project_ids=[proj_id]), + ) + + assert result.id == pool_id + call_args = mock_transport.request.call_args + assert call_args[0][0] == "PATCH" + assert call_args[0][1] == f"/api/v2/agent-pools/{pool_id}" + body = call_args[1]["json_body"]["data"] + assert body["type"] == "agent-pools" + assert body["id"] == pool_id + proj_data = body["relationships"]["allowed-projects"]["data"] + assert proj_data == [{"type": "projects", "id": proj_id}] + + def test_update_allowed_projects_multiple( + self, agent_pools_service, mock_transport + ): + """update_allowed_projects can assign multiple projects at once.""" + pool_id = "apool-123456789abcdef0" + mock_transport.request.return_value.json.return_value = self._pool_response( + pool_id + ) + + agent_pools_service.update_allowed_projects( + pool_id, + AgentPoolAllowedProjectsUpdateOptions(project_ids=["prj-aaa", "prj-bbb"]), + ) + + call_args = mock_transport.request.call_args + body = call_args[1]["json_body"]["data"] + proj_data = body["relationships"]["allowed-projects"]["data"] + assert len(proj_data) == 2 + assert proj_data[0] == {"type": "projects", "id": "prj-aaa"} + assert proj_data[1] == {"type": "projects", "id": "prj-bbb"} + + def test_update_allowed_projects_empty_clears( + self, agent_pools_service, mock_transport + ): + """update_allowed_projects with empty list sends empty array to clear the relationship.""" + pool_id = "apool-123456789abcdef0" + mock_transport.request.return_value.json.return_value = self._pool_response( + pool_id + ) + + agent_pools_service.update_allowed_projects( + pool_id, + AgentPoolAllowedProjectsUpdateOptions(project_ids=[]), + ) + + call_args = mock_transport.request.call_args + body = call_args[1]["json_body"]["data"] + assert "allowed-projects" in body["relationships"] + assert body["relationships"]["allowed-projects"]["data"] == [] + + def test_update_excluded_workspaces(self, agent_pools_service, mock_transport): + """update_excluded_workspaces PATCHes excluded-workspaces relationship.""" + pool_id = "apool-123456789abcdef0" + ws_id = "ws-aaaaaaaaaaaaaaa1" + mock_transport.request.return_value.json.return_value = self._pool_response( + pool_id + ) + + result = agent_pools_service.update_excluded_workspaces( + pool_id, + AgentPoolExcludedWorkspacesUpdateOptions(workspace_ids=[ws_id]), + ) + + assert result.id == pool_id + call_args = mock_transport.request.call_args + assert call_args[0][0] == "PATCH" + assert call_args[0][1] == f"/api/v2/agent-pools/{pool_id}" + body = call_args[1]["json_body"]["data"] + ws_data = body["relationships"]["excluded-workspaces"]["data"] + assert ws_data == [{"type": "workspaces", "id": ws_id}] + + def test_update_excluded_workspaces_empty_clears( + self, agent_pools_service, mock_transport + ): + """update_excluded_workspaces with empty list sends empty array to clear the relationship.""" + pool_id = "apool-123456789abcdef0" + mock_transport.request.return_value.json.return_value = self._pool_response( + pool_id + ) + + agent_pools_service.update_excluded_workspaces( + pool_id, + AgentPoolExcludedWorkspacesUpdateOptions(workspace_ids=[]), + ) + + call_args = mock_transport.request.call_args + body = call_args[1]["json_body"]["data"] + assert "excluded-workspaces" in body["relationships"] + assert body["relationships"]["excluded-workspaces"]["data"] == [] + + def test_update_allowed_projects_invalid_pool_id( + self, agent_pools_service, mock_transport + ): + """update_allowed_projects raises ValueError for invalid pool ID.""" + with pytest.raises(ValueError, match="Agent pool ID"): + agent_pools_service.update_allowed_projects( + "", + AgentPoolAllowedProjectsUpdateOptions(project_ids=["prj-aaa"]), + ) + + def test_update_allowed_workspaces_invalid_pool_id( + self, agent_pools_service, mock_transport + ): + """update_allowed_workspaces raises ValueError for invalid pool ID.""" + with pytest.raises(ValueError, match="Agent pool ID"): + agent_pools_service.update_allowed_workspaces( + "", + AgentPoolAllowedWorkspacesUpdateOptions(workspace_ids=["ws-aaa"]), + ) + + def test_update_excluded_workspaces_invalid_pool_id( + self, agent_pools_service, mock_transport + ): + """update_excluded_workspaces raises ValueError for invalid pool ID.""" + with pytest.raises(ValueError, match="Agent pool ID"): + agent_pools_service.update_excluded_workspaces( + "", + AgentPoolExcludedWorkspacesUpdateOptions(workspace_ids=["ws-aaa"]), + ) + class TestAgentTokenOperations: """Test agent token operations""" From 56562d792c9f4692ceaa942ad28999c499e37ebb Mon Sep 17 00:00:00 2001 From: Tanya Singh Date: Fri, 22 May 2026 17:12:03 +0530 Subject: [PATCH 2/2] Add agent pool project relationship updates --- examples/agent_pool.py | 2 +- src/pytfe/models/agent.py | 10 ++-------- src/pytfe/resources/agent_pools.py | 1 - tests/units/test_agent_pools.py | 27 ++++----------------------- 4 files changed, 7 insertions(+), 33 deletions(-) diff --git a/examples/agent_pool.py b/examples/agent_pool.py index 1d5c37c6..be3e3f83 100644 --- a/examples/agent_pool.py +++ b/examples/agent_pool.py @@ -7,7 +7,7 @@ 1. Agent Pool CRUD operations (Create, Read, Update, Delete) 2. Agent token creation and management 3. Workspace assignment using assign_to_workspaces and remove_from_workspaces -4. Project assignment using update_allowed_projects (Go SDK parity) +4. Project assignment using update_allowed_projects 5. Dedicated relationship update methods: update_allowed_workspaces, update_allowed_projects, update_excluded_workspaces 6. Proper error handling diff --git a/src/pytfe/models/agent.py b/src/pytfe/models/agent.py index c9761203..ac8add4b 100644 --- a/src/pytfe/models/agent.py +++ b/src/pytfe/models/agent.py @@ -142,19 +142,13 @@ class AgentPoolRemoveFromWorkspacesOptions(BaseModel): class AgentPoolAllowedWorkspacesUpdateOptions(BaseModel): - """Options for updating the allowed-workspaces relationship on an agent pool. - - Supports full replacement including clearing (pass an empty list). - """ + """Options for updating the allowed-workspaces relationship on an agent pool.""" workspace_ids: list[str] = Field(default_factory=list) class AgentPoolAllowedProjectsUpdateOptions(BaseModel): - """Options for updating the allowed-projects relationship on an agent pool. - - Supports full replacement including clearing (pass an empty list). - """ + """Options for updating the allowed-projects relationship on an agent pool.""" project_ids: list[str] = Field(default_factory=list) diff --git a/src/pytfe/resources/agent_pools.py b/src/pytfe/resources/agent_pools.py index 344d3628..325eaa57 100644 --- a/src/pytfe/resources/agent_pools.py +++ b/src/pytfe/resources/agent_pools.py @@ -586,7 +586,6 @@ def _patch_relationship( Always sends the relationship array in the payload even when empty, which allows callers to clear an existing relationship by passing an empty list. - This mirrors the Go SDK updateArrayAttribute helper. """ if not valid_string_id(agent_pool_id): raise ValueError("Agent pool ID is required and must be valid") diff --git a/tests/units/test_agent_pools.py b/tests/units/test_agent_pools.py index e0bfed65..24f0cc37 100644 --- a/tests/units/test_agent_pools.py +++ b/tests/units/test_agent_pools.py @@ -10,9 +10,6 @@ 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 """ from unittest.mock import Mock @@ -95,7 +92,7 @@ def test_agent_pool_create_options(self): ) def test_agent_pool_create_options_workspace_ids(self): - """Test AgentPoolCreateOptions with allowed/excluded workspace IDs (bug fix)""" + """Test AgentPoolCreateOptions with allowed/excluded workspace IDs""" options = AgentPoolCreateOptions( name="scoped-pool", organization_scoped=False, @@ -115,7 +112,7 @@ def test_agent_pool_create_options_project_ids(self): assert options.allowed_project_ids == ["prj-aaa", "prj-bbb"] def test_agent_pool_update_options_workspace_ids(self): - """Test AgentPoolUpdateOptions with allowed/excluded workspace IDs (bug fix)""" + """Test AgentPoolUpdateOptions with allowed/excluded workspace IDs""" options = AgentPoolUpdateOptions( allowed_workspace_ids=["ws-aaa"], excluded_workspace_ids=["ws-bbb"], @@ -196,7 +193,6 @@ def test_list_agent_pools(self, agent_pools_service, mock_transport): assert agent_pools[0].name == "test-pool-1" assert agent_pools[0].agent_count == 2 - # Verify API call mock_transport.request.assert_called_once() call_args = mock_transport.request.call_args assert "organizations/test-org/agent-pools" in call_args[0][1] @@ -250,7 +246,6 @@ def test_create_agent_pool(self, agent_pools_service, mock_transport): assert agent_pool.name == "new-pool" assert agent_pool.organization_scoped is True - # Verify API call mock_transport.request.assert_called_once() call_args = mock_transport.request.call_args assert call_args[0][0] == "POST" @@ -280,7 +275,6 @@ def test_read_agent_pool(self, agent_pools_service, mock_transport): assert agent_pool.organization_scoped is False assert agent_pool.agent_count == 3 - # Verify API call mock_transport.request.assert_called_once() call_args = mock_transport.request.call_args assert call_args[0][0] == "GET" @@ -310,7 +304,6 @@ def test_update_agent_pool(self, agent_pools_service, mock_transport): assert agent_pool.name == "updated-pool" assert agent_pool.organization_scoped is False - # Verify API call mock_transport.request.assert_called_once() call_args = mock_transport.request.call_args assert call_args[0][0] == "PATCH" @@ -320,7 +313,6 @@ def test_delete_agent_pool(self, agent_pools_service, mock_transport): """Test deleting an agent pool""" agent_pools_service.delete("apool-123456789abcdef0") - # Verify API call mock_transport.request.assert_called_once() call_args = mock_transport.request.call_args assert call_args[0][0] == "DELETE" @@ -359,11 +351,8 @@ def test_assign_to_workspaces(self, agent_pools_service, mock_transport): 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 @@ -374,7 +363,6 @@ def test_assign_to_workspaces(self, agent_pools_service, mock_transport): 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" @@ -404,11 +392,8 @@ def test_remove_from_workspaces(self, agent_pools_service, mock_transport): 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 @@ -481,7 +466,7 @@ def test_update_agent_pool_with_project_ids( class TestDedicatedRelationshipUpdateMethods: - """Test the Go SDK-parity dedicated update methods for relationships.""" + """Test the dedicated update methods for relationships.""" @pytest.fixture def mock_transport(self): @@ -785,7 +770,6 @@ def test_read_agent_token(self, agent_tokens_service, mock_transport): assert token.id == "at-123456789abcdef0" assert token.description == "Existing token" - # Verify API call mock_transport.request.assert_called_once() call_args = mock_transport.request.call_args assert call_args[0][0] == "GET" @@ -795,7 +779,6 @@ def test_delete_agent_token(self, agent_tokens_service, mock_transport): """Test deleting an agent token""" agent_tokens_service.delete("at-123456789abcdef0") - # Verify API call mock_transport.request.assert_called_once() call_args = mock_transport.request.call_args assert call_args[0][0] == "DELETE" @@ -829,9 +812,7 @@ def test_validation_error(self, agent_pools_service, mock_transport): """Test handling of ValidationError errors""" mock_transport.request.side_effect = ValidationError("Invalid agent pool name") - options = AgentPoolCreateOptions( - name="valid-name" - ) # Use valid name to avoid ValueError + options = AgentPoolCreateOptions(name="valid-name") with pytest.raises(ValidationError): agent_pools_service.create("test-org", options)