From 111012ed1fe84a44727761c0e1612ce9cb01a36e Mon Sep 17 00:00:00 2001 From: Sivaselvan32 Date: Sat, 23 May 2026 22:04:24 +0530 Subject: [PATCH 1/5] refactor(agent-pool): Updated models to include project_ids, workspace_ids for allowed and excluded, removed allowed_workspace_policy in models, updated relationships for AgentPool, updated listOptions with parameters --- src/pytfe/models/agent.py | 139 +++++++++++++++++++++++++------------- 1 file changed, 93 insertions(+), 46 deletions(-) diff --git a/src/pytfe/models/agent.py b/src/pytfe/models/agent.py index d0751ea1..76c02f3c 100644 --- a/src/pytfe/models/agent.py +++ b/src/pytfe/models/agent.py @@ -11,9 +11,20 @@ from datetime import datetime from enum import Enum -from typing import Any +from typing import TYPE_CHECKING -from pydantic import BaseModel, Field +from pydantic import BaseModel, ConfigDict, Field, model_validator + +from ..errors import ( + InvalidNameError, + RequiredNameError, +) +from ..utils import valid_string, valid_string_id +from .organization import Organization +from .workspace import Workspace + +if TYPE_CHECKING: + from .project import Project class AgentStatus(str, Enum): @@ -24,13 +35,6 @@ class AgentStatus(str, Enum): UNKNOWN = "unknown" -class AgentPoolAllowedWorkspacePolicy(str, Enum): - """Agent pool allowed workspace policy enumeration.""" - - ALL_WORKSPACES = "all-workspaces" - SPECIFIC_WORKSPACES = "specific-workspaces" - - class Agent(BaseModel): """Agent represents a Terraform Enterprise agent.""" @@ -48,72 +52,112 @@ class Agent(BaseModel): class AgentPool(BaseModel): """Agent Pool represents a Terraform Enterprise agent pool.""" + model_config = ConfigDict(populate_by_name=True, validate_by_name=True) + id: str - name: str | None = None - created_at: datetime | None = None - organization_scoped: bool | None = None - allowed_workspace_policy: AgentPoolAllowedWorkspacePolicy | None = None - agent_count: int = 0 + name: str | None = Field(default=None, alias="name") + created_at: datetime | None = Field(default=None, alias="created-at") + organization_scoped: bool | None = Field(default=None, alias="organization-scoped") + agent_count: int | None = Field(default=None, alias="agent-count") # Relations - organization: Any | None = None # Organization type from main types - workspaces: list[Any] = Field(default_factory=list) # Workspace types + organization: Organization | None = Field(default=None, alias="organization") + workspaces: list[Workspace] = Field(default_factory=list, alias="workspaces") agents: list[Agent] = Field(default_factory=list) + allowed_workspaces: list[Workspace] = Field( + default_factory=list, alias="allowed-workspaces" + ) + excluded_workspaces: list[Workspace] = Field( + default_factory=list, alias="excluded-workspaces" + ) + allowed_projects: list[Project] = Field( + default_factory=list, alias="allowed-projects" + ) -# Agent Pool Options +class AgentPoolIncludeOpt(str, Enum): + AGENT_POOL_WORKSPACES = "workspaces" class AgentPoolListOptions(BaseModel): """Options for listing agent pools.""" - # Pagination options - page_number: int | None = None - page_size: int | None = None + model_config = ConfigDict(populate_by_name=True, validate_by_name=True) + + page_size: int | None = Field(default=None, alias="page[size]") # Optional: Include related resources - include: list[str] | None = None - # Optional: Filter by allowed workspace policy - allowed_workspace_policy: AgentPoolAllowedWorkspacePolicy | None = None + include: list[AgentPoolIncludeOpt] | None = Field(default=None, alias="include") + query: str | None = Field(default=None, alias="q") + allowed_workspace_name: str | None = Field( + default=None, alias="filter[allowed_workspaces][name]" + ) + allowed_project_name: str | None = Field( + default=None, alias="filter[allowed_projects][name]" + ) + sort: str | None = Field(default=None, alias="sort") class AgentPoolCreateOptions(BaseModel): """Options for creating an agent pool.""" + model_config = ConfigDict(populate_by_name=True, validate_by_name=True) + # Required: A name to identify the agent pool - name: str - # Optional: Whether the agent pool is organization scoped - 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) + name: str = Field(alias="name") + organization_scoped: bool | None = Field(default=None, alias="organization-scoped") + allowed_workspace_ids: list[str] | None = Field( + default=None, alias="allowed-workspaces" + ) + excluded_workspace_ids: list[str] | None = Field( + default=None, alias="excluded-workspaces" + ) + allowed_project_ids: list[str] | None = Field( + default=None, alias="allowed-projects" + ) + + @model_validator(mode="after") + def valid(self) -> AgentPoolCreateOptions: + """Validate the options for creating an agent pool.""" + if not valid_string(self.name): + raise RequiredNameError() + if not valid_string_id(self.name): + raise InvalidNameError() + + return self class AgentPoolUpdateOptions(BaseModel): """Options for updating an agent pool.""" + model_config = ConfigDict(populate_by_name=True, validate_by_name=True) + # Optional: A name to identify the agent pool - name: str | None = None - # Optional: Whether the agent pool is organization scoped - 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) + name: str | None = Field(default=None, alias="name") + organization_scoped: bool | None = Field(default=None, alias="organization-scoped") + allowed_workspace_ids: list[str] | None = Field( + default=None, alias="allowed-workspaces" + ) + excluded_workspace_ids: list[str] | None = Field( + default=None, alias="excluded-workspaces" + ) + allowed_project_ids: list[str] | None = Field( + default=None, alias="allowed-projects" + ) + + @model_validator(mode="after") + def valid(self) -> AgentPoolUpdateOptions: + """Validate the options for updating an agent pool.""" + if self.name is not None and not valid_string_id(self.name): + raise InvalidNameError() + + return self class AgentPoolReadOptions(BaseModel): """Options for reading an agent pool.""" # Optional: Include related resources - include: list[str] | None = None - - -# Agent Pool Workspace Assignment Options + include: list[AgentPoolIncludeOpt] | None = Field(default=None, alias="include") class AgentPoolAssignToWorkspacesOptions(BaseModel): @@ -128,7 +172,10 @@ class AgentPoolRemoveFromWorkspacesOptions(BaseModel): workspace_ids: list[str] = Field(default_factory=list) -# Agent Options +class AgentPoolAssignToProjectsOptions(BaseModel): + """Options for assigning an agent pool to projects.""" + + project_ids: list[str] = Field(default_factory=list) class AgentListOptions(BaseModel): From 86e96f75badfeeeffeaad9a09223807aa7a08cc9 Mon Sep 17 00:00:00 2001 From: Sivaselvan32 Date: Sat, 23 May 2026 22:09:31 +0530 Subject: [PATCH 2/5] refactor(agent-pool): Added assign_to_project method to the agent-pool resource to add projects, added _parse_agent_pool_from method, removed allowed_workspace_policy in the methods, added valid options in create/update in the models --- src/pytfe/resources/agent_pools.py | 400 +++++++++++------------------ 1 file changed, 145 insertions(+), 255 deletions(-) diff --git a/src/pytfe/resources/agent_pools.py b/src/pytfe/resources/agent_pools.py index 47ffc8bf..726368c7 100644 --- a/src/pytfe/resources/agent_pools.py +++ b/src/pytfe/resources/agent_pools.py @@ -10,11 +10,24 @@ from __future__ import annotations from collections.abc import Iterator -from typing import Any, cast - +from typing import Any + +from pytfe.models.organization import Organization +from pytfe.models.project import Project +from pytfe.models.workspace import Workspace + +from ..errors import ( + InvalidAgentPoolIDError, + InvalidOrgError, + InvalidProjectIDError, + InvalidWorkspaceIDError, + RequiredProjectError, + RequiredWorkspaceError, +) from ..models.agent import ( + Agent, AgentPool, - AgentPoolAllowedWorkspacePolicy, + AgentPoolAssignToProjectsOptions, AgentPoolAssignToWorkspacesOptions, AgentPoolCreateOptions, AgentPoolListOptions, @@ -22,91 +35,10 @@ AgentPoolRemoveFromWorkspacesOptions, AgentPoolUpdateOptions, ) -from ..utils import valid_string, valid_string_id +from ..utils import valid_string_id from ._base import _Service -def valid_agent_pool_name(name: str) -> bool: - """Validate agent pool name format.""" - if not valid_string(name): - return False - # Agent pool names must be between 1 and 90 characters - # and can contain letters, numbers, spaces, hyphens, and underscores - if len(name) > 90: - return False - return True - - -def validate_agent_pool_create_options(organization: str, name: str) -> None: - """Validate agent pool creation parameters.""" - if not valid_string(organization): - raise ValueError("Organization name is required and must be valid") - - if not valid_string(name): - raise ValueError("Agent pool name is required") - - if not valid_agent_pool_name(name): - raise ValueError("Agent pool name contains invalid characters or is too long") - - -def validate_agent_pool_update_options( - agent_pool_id: str, name: str | None = None -) -> None: - """Validate agent pool update parameters.""" - if not valid_string_id(agent_pool_id): - raise ValueError("Agent pool ID is required and must be valid") - - if name is not None: - if not valid_string(name): - raise ValueError("Agent pool name must be a valid string") - if not valid_agent_pool_name(name): - raise ValueError( - "Agent pool name contains invalid characters or is too long" - ) - - -def _safe_str(value: Any, default: str = "") -> str: - """Safely convert a value to string with optional default.""" - if value is None: - return default - return str(value) - - -def _safe_int(value: Any, default: int = 0) -> int: - """Safely convert a value to an integer.""" - if value is None: - return default - if isinstance(value, int): - return value - try: - return int(value) - except (ValueError, TypeError): - return default - - -def _safe_bool(value: Any) -> bool | None: - """Safely convert a value to a boolean.""" - if value is None: - return None - if isinstance(value, bool): - return value - if isinstance(value, str): - return value.lower() in ("true", "1", "yes", "on") - return bool(value) - - -def _safe_workspace_policy(value: Any) -> AgentPoolAllowedWorkspacePolicy | None: - """Safely convert a value to an AgentPoolAllowedWorkspacePolicy enum.""" - if value is None: - return None - if isinstance(value, AgentPoolAllowedWorkspacePolicy): - return value - try: - return AgentPoolAllowedWorkspacePolicy(str(value)) - except (ValueError, TypeError): - return None - - class AgentPools(_Service): """Agent Pools service for managing Terraform Enterprise agent pools.""" @@ -126,58 +58,29 @@ def list( ValueError: If organization name is invalid TFEError: If API request fails """ - if not valid_string(organization): - raise ValueError("Organization name is required and must be valid") + if not valid_string_id(organization): + raise InvalidOrgError() path = f"/api/v2/organizations/{organization}/agent-pools" params: dict[str, str | int] = {} if options: - if options.page_number is not None: - params["page[number]"] = options.page_number if options.page_size is not None: params["page[size]"] = options.page_size if options.include: params["include"] = ",".join(options.include) - if options.allowed_workspace_policy: - params["filter[allowed_workspace_policy]"] = ( - options.allowed_workspace_policy.value + if options.query: + params["q"] = options.query + if options.allowed_workspace_name: + params["filter[allowed_workspaces][name]"] = ( + options.allowed_workspace_name ) - - items_iter = self._list(path, params=params) - - for item in items_iter: - # Extract agent pool data from API response - attr = item.get("attributes", {}) or {} - relationships = item.get("relationships", {}) or {} - - # Note: organization and workspace relationships available but not currently used - - # Extract agents from relationships - agents_data = relationships.get("agents", {}).get("data", []) - agent_count = ( - len(agents_data) if agents_data else attr.get("agent-count", 0) - ) - - agent_pool_data = { - "id": _safe_str(item.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": agent_count, - } - - yield 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"]), - ) + if options.allowed_project_name: + params["filter[allowed_projects][name]"] = options.allowed_project_name + if options.sort: + params["sort"] = options.sort + for item in self._list(path, params=params): + yield self._parse_agent_pool_from(item) def create(self, organization: str, options: AgentPoolCreateOptions) -> AgentPool: """Create a new agent pool in an organization. @@ -193,7 +96,8 @@ def create(self, organization: str, options: AgentPoolCreateOptions) -> AgentPoo ValueError: If parameters are invalid TFEError: If API request fails """ - validate_agent_pool_create_options(organization, options.name) + if not valid_string_id(organization): + raise InvalidOrgError() path = f"/api/v2/organizations/{organization}/agent-pools" attributes: dict[str, Any] = {"name": options.name} @@ -201,11 +105,6 @@ def create(self, organization: str, options: AgentPoolCreateOptions) -> AgentPoo if options.organization_scoped is not None: attributes["organization-scoped"] = options.organization_scoped - if options.allowed_workspace_policy is not None: - attributes["allowed-workspace-policy"] = ( - options.allowed_workspace_policy.value - ) - relationships: dict[str, Any] = {} if options.allowed_workspace_ids: relationships["allowed-workspaces"] = { @@ -221,6 +120,13 @@ def create(self, organization: str, options: AgentPoolCreateOptions) -> AgentPoo for ws_id in options.excluded_workspace_ids ] } + if options.allowed_project_ids: + relationships["allowed-projects"] = { + "data": [ + {"type": "projects", "id": proj_id} + for proj_id in options.allowed_project_ids + ] + } payload: dict[str, Any] = { "data": {"type": "agent-pools", "attributes": attributes} @@ -231,27 +137,7 @@ def create(self, organization: str, options: AgentPoolCreateOptions) -> AgentPoo response = self.t.request("POST", 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"]), - ) + return self._parse_agent_pool_from(data) def read( self, agent_pool_id: str, options: AgentPoolReadOptions | None = None @@ -270,7 +156,7 @@ def read( TFEError: If API request fails """ if not valid_string_id(agent_pool_id): - raise ValueError("Agent pool ID is required and must be valid") + raise InvalidAgentPoolIDError() path = f"/api/v2/agent-pools/{agent_pool_id}" params: dict[str, str] = {} @@ -285,33 +171,7 @@ def read( data = response.json()["data"] - # Extract agent pool data from response - attr = data.get("attributes", {}) or {} - relationships = data.get("relationships", {}) or {} - - # Extract agents count - agents_data = relationships.get("agents", {}).get("data", []) - agent_count = len(agents_data) if agents_data else attr.get("agent-count", 0) - - 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": agent_count, - } - - 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"]), - ) + return self._parse_agent_pool_from(data) def update(self, agent_pool_id: str, options: AgentPoolUpdateOptions) -> AgentPool: """Update an agent pool's properties. @@ -327,7 +187,9 @@ def update(self, agent_pool_id: str, options: AgentPoolUpdateOptions) -> AgentPo ValueError: If parameters are invalid TFEError: If API request fails """ - validate_agent_pool_update_options(agent_pool_id, options.name) + + if not valid_string_id(agent_pool_id): + raise InvalidAgentPoolIDError() path = f"/api/v2/agent-pools/{agent_pool_id}" attributes: dict[str, Any] = {} @@ -338,11 +200,6 @@ def update(self, agent_pool_id: str, options: AgentPoolUpdateOptions) -> AgentPo if options.organization_scoped is not None: attributes["organization-scoped"] = options.organization_scoped - if options.allowed_workspace_policy is not None: - attributes["allowed-workspace-policy"] = ( - options.allowed_workspace_policy.value - ) - relationships: dict[str, Any] = {} if options.allowed_workspace_ids: relationships["allowed-workspaces"] = { @@ -358,6 +215,13 @@ def update(self, agent_pool_id: str, options: AgentPoolUpdateOptions) -> AgentPo for ws_id in options.excluded_workspace_ids ] } + if options.allowed_project_ids: + relationships["allowed-projects"] = { + "data": [ + {"type": "projects", "id": proj_id} + for proj_id in options.allowed_project_ids + ] + } payload: dict[str, Any] = { "data": { @@ -372,27 +236,7 @@ def update(self, agent_pool_id: str, options: AgentPoolUpdateOptions) -> AgentPo 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"]), - ) + return self._parse_agent_pool_from(data) def delete(self, agent_pool_id: str) -> None: """Delete an agent pool. @@ -405,7 +249,7 @@ def delete(self, agent_pool_id: str) -> None: TFEError: If API request fails """ if not valid_string_id(agent_pool_id): - raise ValueError("Agent pool ID is required and must be valid") + raise InvalidAgentPoolIDError() path = f"/api/v2/agent-pools/{agent_pool_id}" self.t.request("DELETE", path) @@ -431,14 +275,14 @@ def assign_to_workspaces( TFEError: If API request fails """ if not valid_string_id(agent_pool_id): - raise ValueError("Agent pool ID is required and must be valid") + raise InvalidAgentPoolIDError() if not options.workspace_ids: - raise ValueError("At least one workspace ID is required") + raise RequiredWorkspaceError() for workspace_id in options.workspace_ids: if not valid_string_id(workspace_id): - raise ValueError(f"Invalid workspace ID: {workspace_id}") + raise InvalidWorkspaceIDError(f"Invalid workspace ID: {workspace_id}") path = f"/api/v2/agent-pools/{agent_pool_id}" payload: dict[str, Any] = { @@ -459,27 +303,7 @@ def assign_to_workspaces( 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"]), - ) + return self._parse_agent_pool_from(data) def remove_from_workspaces( self, agent_pool_id: str, options: AgentPoolRemoveFromWorkspacesOptions @@ -503,14 +327,14 @@ def remove_from_workspaces( TFEError: If API request fails """ if not valid_string_id(agent_pool_id): - raise ValueError("Agent pool ID is required and must be valid") + raise InvalidAgentPoolIDError() if not options.workspace_ids: - raise ValueError("At least one workspace ID is required") + raise RequiredWorkspaceError() for workspace_id in options.workspace_ids: if not valid_string_id(workspace_id): - raise ValueError(f"Invalid workspace ID: {workspace_id}") + raise InvalidWorkspaceIDError(f"Invalid workspace ID: {workspace_id}") path = f"/api/v2/agent-pools/{agent_pool_id}" payload: dict[str, Any] = { @@ -531,24 +355,90 @@ def remove_from_workspaces( 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 self._parse_agent_pool_from(data) + + def assign_to_projects( + self, agent_pool_id: str, options: AgentPoolAssignToProjectsOptions + ) -> AgentPool: + """Assign an agent pool to projects by updating the allowed-projects + relationship via PATCH /agent-pools/:id. + + The provided project IDs become the new complete list of allowed + projects for this pool (full replacement, not append). + + Args: + agent_pool_id: Agent pool ID + options: Assignment options containing project IDs + """ + if not valid_string_id(agent_pool_id): + raise InvalidAgentPoolIDError() + + if not options.project_ids: + raise RequiredProjectError() + + for project_id in options.project_ids: + if not valid_string_id(project_id): + raise InvalidProjectIDError(f"Invalid project ID: {project_id}") + + path = f"/api/v2/agent-pools/{agent_pool_id}" + payload: dict[str, Any] = { + "data": { + "type": "agent-pools", + "id": agent_pool_id, + "attributes": {}, + "relationships": { + "allowed-projects": { + "data": [ + {"type": "projects", "id": project_id} + for project_id in options.project_ids + ] + } + }, + } } + response = self.t.request("PATCH", path, json_body=payload) + data = response.json()["data"] + + return self._parse_agent_pool_from(data) + + def _parse_agent_pool_from(self, data: dict[str, Any]) -> AgentPool: + """Helper method to parse agent pool data from API response.""" + attr = data.get("attributes", {}) + relationships = data.get("relationships", {}) + attr["id"] = data.get("id") + + # Extract agents count + agents_data = relationships.get("agents", {}).get("data", []) + attr["agents"] = [Agent(id=agent["id"]) for agent in agents_data] + + org_data = relationships.get("organization", {}).get("data") + attr["organization"] = Organization(id=org_data["id"]) if org_data else None - 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"]), + workspaces_data = relationships.get("workspaces", {}).get("data", []) + attr["workspaces"] = [ + Workspace.model_validate({"id": ws["id"]}) for ws in workspaces_data + ] + + allowed_workspaces_data = relationships.get("allowed-workspaces", {}).get( + "data", [] ) + attr["allowed_workspaces"] = [ + Workspace.model_validate({"id": ws["id"]}) for ws in allowed_workspaces_data + ] + + excluded_workspaces_data = relationships.get("excluded-workspaces", {}).get( + "data", [] + ) + attr["excluded_workspaces"] = [ + Workspace.model_validate({"id": ws["id"]}) + for ws in excluded_workspaces_data + ] + + allowed_projects_data = relationships.get("allowed-projects", {}).get( + "data", [] + ) + attr["allowed_projects"] = [ + Project.model_validate({"id": proj["id"]}) for proj in allowed_projects_data + ] + + return AgentPool.model_validate(attr) From 4597aa54029f4ee1769e08743ab6660d07630233 Mon Sep 17 00:00:00 2001 From: Sivaselvan32 Date: Sat, 23 May 2026 22:11:15 +0530 Subject: [PATCH 3/5] refactor(agent-pool): added agent-pool errors, added models in init, updated examples and testcases --- examples/agent_pool.py | 15 +++++++++++--- src/pytfe/errors.py | 15 ++++++++++++++ src/pytfe/models/__init__.py | 20 +++++++++++++++++-- tests/units/test_agent_pools.py | 35 --------------------------------- 4 files changed, 45 insertions(+), 40 deletions(-) diff --git a/examples/agent_pool.py b/examples/agent_pool.py index bcb04ae3..880a86d7 100644 --- a/examples/agent_pool.py +++ b/examples/agent_pool.py @@ -27,13 +27,13 @@ from pytfe import TFEClient, TFEConfig from pytfe.errors import NotFound from pytfe.models import ( - AgentPoolAllowedWorkspacePolicy, AgentPoolAssignToWorkspacesOptions, AgentPoolCreateOptions, AgentPoolListOptions, AgentPoolRemoveFromWorkspacesOptions, AgentPoolUpdateOptions, AgentTokenCreateOptions, + AgentPoolAssignToProjectsOptions, ) @@ -46,6 +46,9 @@ 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") @@ -81,7 +84,6 @@ def main(): create_options = AgentPoolCreateOptions( name=unique_name, organization_scoped=True, # Optional parameter - allowed_workspace_policy=AgentPoolAllowedWorkspacePolicy.ALL_WORKSPACES, # Optional ) new_pool = client.agent_pools.create(org, create_options) @@ -92,7 +94,6 @@ def main(): pool_details = client.agent_pools.read(new_pool.id) print(f"Name: {pool_details.name}") print(f"Organization Scoped: {pool_details.organization_scoped}") - print(f"Policy: {pool_details.allowed_workspace_policy}") print(f"Agent Count: {pool_details.agent_count}") # Example 4: Update the agent pool @@ -125,6 +126,14 @@ def main(): else: print("\n Skipping workspace assignment (set TFE_WORKSPACE_ID to test)") + if project_id: + print("\n Assigning project to agent pool...") + updated_pool = client.agent_pools.assign_to_projects( + new_pool.id, + AgentPoolAssignToProjectsOptions(project_ids=[project_id]), + ) + print(f" Assigned project {project_id} to pool {updated_pool.name}") + # Example 6: Create an agent token print("\n Creating agent token...") token_options = AgentTokenCreateOptions( diff --git a/src/pytfe/errors.py b/src/pytfe/errors.py index 45615818..dbf4ba38 100644 --- a/src/pytfe/errors.py +++ b/src/pytfe/errors.py @@ -702,3 +702,18 @@ class InvalidTokenIDError(InvalidValues): def __init__(self, message: str = "invalid value for token ID"): super().__init__(message) + + +# Agent Pool errors +class InvalidAgentPoolIDError(InvalidValues): + """Raised when an invalid agent pool ID is provided.""" + + def __init__(self, message: str = "invalid value for agent pool ID"): + super().__init__(message) + + +class RequiredProjectError(RequiredFieldMissing): + """Raised when a required project field is missing.""" + + def __init__(self, message: str = "project is required"): + super().__init__(message) diff --git a/src/pytfe/models/__init__.py b/src/pytfe/models/__init__.py index 15e0c4d6..bdbf42d4 100644 --- a/src/pytfe/models/__init__.py +++ b/src/pytfe/models/__init__.py @@ -8,8 +8,8 @@ Agent, AgentListOptions, AgentPool, - AgentPoolAllowedWorkspacePolicy, AgentPoolAssignToWorkspacesOptions, + AgentPoolAssignToProjectsOptions, AgentPoolCreateOptions, AgentPoolListOptions, AgentPoolReadOptions, @@ -509,8 +509,8 @@ # Agent & pools "Agent", "AgentPool", - "AgentPoolAllowedWorkspacePolicy", "AgentPoolAssignToWorkspacesOptions", + "AgentPoolAssignToProjectsOptions", "AgentPoolCreateOptions", "AgentPoolListOptions", "AgentPoolReadOptions", @@ -867,3 +867,19 @@ "TaskStage": TaskStage, }, ) +AgentPool.model_rebuild( + raise_errors=False, + _types_namespace={"Project": Project}, +) +Project.model_rebuild( + raise_errors=False, + _types_namespace={"AgentPool": AgentPool}, +) +RunTask.model_rebuild( + raise_errors=False, + _types_namespace={"WorkspaceRunTask": WorkspaceRunTask, "AgentPool": AgentPool}, +) +Workspace.model_rebuild( + raise_errors=False, + _types_namespace={"AgentPool": AgentPool, "Run": Run, "TaskStage": TaskStage}, +) diff --git a/tests/units/test_agent_pools.py b/tests/units/test_agent_pools.py index f797f612..9b55fe99 100644 --- a/tests/units/test_agent_pools.py +++ b/tests/units/test_agent_pools.py @@ -22,7 +22,6 @@ from pytfe.errors import AuthError, NotFound, ValidationError from pytfe.models.agent import ( AgentPool, - AgentPoolAllowedWorkspacePolicy, AgentPoolAssignToWorkspacesOptions, AgentPoolCreateOptions, AgentPoolListOptions, @@ -42,54 +41,23 @@ def test_agent_pool_model_basic(self): name="test-pool", created_at="2023-01-01T00:00:00Z", organization_scoped=True, - allowed_workspace_policy=AgentPoolAllowedWorkspacePolicy.ALL_WORKSPACES, agent_count=0, ) assert agent_pool.id == "apool-123456789abcdef0" assert agent_pool.name == "test-pool" assert agent_pool.organization_scoped is True - assert ( - agent_pool.allowed_workspace_policy - == AgentPoolAllowedWorkspacePolicy.ALL_WORKSPACES - ) assert agent_pool.agent_count == 0 - def test_agent_pool_allowed_workspace_policy_enum(self): - """Test AgentPoolAllowedWorkspacePolicy enum values""" - assert AgentPoolAllowedWorkspacePolicy.ALL_WORKSPACES == "all-workspaces" - assert ( - AgentPoolAllowedWorkspacePolicy.SPECIFIC_WORKSPACES == "specific-workspaces" - ) - - agent_pool = AgentPool( - id="apool-123456789abcdef0", - name="test-pool", - created_at="2023-01-01T00:00:00Z", - organization_scoped=False, - allowed_workspace_policy=AgentPoolAllowedWorkspacePolicy.SPECIFIC_WORKSPACES, - agent_count=3, - ) - - assert ( - agent_pool.allowed_workspace_policy - == AgentPoolAllowedWorkspacePolicy.SPECIFIC_WORKSPACES - ) - def test_agent_pool_create_options(self): """Test AgentPoolCreateOptions model""" options = AgentPoolCreateOptions( name="test-pool", organization_scoped=True, - allowed_workspace_policy=AgentPoolAllowedWorkspacePolicy.SPECIFIC_WORKSPACES, ) assert options.name == "test-pool" assert options.organization_scoped is True - assert ( - options.allowed_workspace_policy - == AgentPoolAllowedWorkspacePolicy.SPECIFIC_WORKSPACES - ) def test_agent_pool_create_options_workspace_ids(self): """Test AgentPoolCreateOptions with allowed/excluded workspace IDs (bug fix)""" @@ -165,7 +133,6 @@ def test_list_agent_pools_with_options(self, agent_pools_service, mock_transport options = AgentPoolListOptions( page_size=10, - allowed_workspace_policy=AgentPoolAllowedWorkspacePolicy.ALL_WORKSPACES, ) list(agent_pools_service.list("test-org", options)) @@ -176,7 +143,6 @@ def test_list_agent_pools_with_options(self, agent_pools_service, mock_transport params = call_args[1]["params"] assert params["page[number]"] == 1 assert params["page[size]"] == 10 - assert params["filter[allowed_workspace_policy]"] == "all-workspaces" def test_create_agent_pool(self, agent_pools_service, mock_transport): """Test creating an agent pool""" @@ -198,7 +164,6 @@ def test_create_agent_pool(self, agent_pools_service, mock_transport): options = AgentPoolCreateOptions( name="new-pool", organization_scoped=True, - allowed_workspace_policy=AgentPoolAllowedWorkspacePolicy.ALL_WORKSPACES, ) agent_pool = agent_pools_service.create("test-org", options) From 2da9da1081c02e4df3a95f2faecd72a01f43d238 Mon Sep 17 00:00:00 2001 From: Sivaselvan32 Date: Sat, 23 May 2026 22:12:17 +0530 Subject: [PATCH 4/5] Updated models to handle circular import issues --- src/pytfe/models/project.py | 6 +++++- src/pytfe/models/workspace.py | 3 ++- src/pytfe/models/workspace_run_task.py | 2 +- 3 files changed, 8 insertions(+), 3 deletions(-) diff --git a/src/pytfe/models/project.py b/src/pytfe/models/project.py index b67d3c9b..535a0f59 100644 --- a/src/pytfe/models/project.py +++ b/src/pytfe/models/project.py @@ -3,12 +3,16 @@ from __future__ import annotations +from typing import TYPE_CHECKING + from pydantic import BaseModel, ConfigDict, Field -from .agent import AgentPool from .common import TagBinding from .organization import Organization +if TYPE_CHECKING: + from .agent import AgentPool + class Project(BaseModel): model_config = ConfigDict(populate_by_name=True, validate_by_name=True) diff --git a/src/pytfe/models/workspace.py b/src/pytfe/models/workspace.py index ee067f43..79e5113e 100644 --- a/src/pytfe/models/workspace.py +++ b/src/pytfe/models/workspace.py @@ -21,7 +21,6 @@ UnsupportedOperationsError, ) from ..utils import has_tags_regex_defined, is_valid_workspace_name, valid_string -from .agent import AgentPool from .common import EffectiveTagBinding, Tag, TagBinding from .configuration_version import ConfigurationVersion from .data_retention_policy import DataRetentionPolicyChoice @@ -32,6 +31,7 @@ from .variable import Variable if TYPE_CHECKING: + from .agent import AgentPool from .run import Run @@ -525,6 +525,7 @@ class VCSRepoOptions(BaseModel): def _rebuild_workspace_model() -> None: """Rebuild Workspace model to resolve forward references.""" try: + from .agent import AgentPool # noqa: F401 from .run import Run # noqa: F401 from .task_stage import TaskStage # noqa: F401 diff --git a/src/pytfe/models/workspace_run_task.py b/src/pytfe/models/workspace_run_task.py index d92fb328..f775f016 100644 --- a/src/pytfe/models/workspace_run_task.py +++ b/src/pytfe/models/workspace_run_task.py @@ -68,4 +68,4 @@ class WorkspaceRunTaskUpdateOptions(BaseModel): # WorkspaceRunTask is now fully defined; rebuild RunTask so Pydantic can # resolve the forward reference in RunTask.workspace_run_tasks. -RunTask.model_rebuild() +RunTask.model_rebuild(raise_errors=False) From cec7be0413dbc9b728ceb0deefed2a58ac4d8a20 Mon Sep 17 00:00:00 2001 From: Sivaselvan32 Date: Sat, 23 May 2026 22:15:58 +0530 Subject: [PATCH 5/5] fixed fmt and lint --- examples/agent_pool.py | 6 ++---- src/pytfe/models/__init__.py | 2 +- 2 files changed, 3 insertions(+), 5 deletions(-) diff --git a/examples/agent_pool.py b/examples/agent_pool.py index 880a86d7..3fa086ad 100644 --- a/examples/agent_pool.py +++ b/examples/agent_pool.py @@ -27,13 +27,13 @@ from pytfe import TFEClient, TFEConfig from pytfe.errors import NotFound from pytfe.models import ( + AgentPoolAssignToProjectsOptions, AgentPoolAssignToWorkspacesOptions, AgentPoolCreateOptions, AgentPoolListOptions, AgentPoolRemoveFromWorkspacesOptions, AgentPoolUpdateOptions, AgentTokenCreateOptions, - AgentPoolAssignToProjectsOptions, ) @@ -46,9 +46,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 + project_id = os.environ.get("TFE_PROJECT_ID") # optional, for project assignment if not token: print("TFE_TOKEN environment variable is required") diff --git a/src/pytfe/models/__init__.py b/src/pytfe/models/__init__.py index bdbf42d4..04c9aaa7 100644 --- a/src/pytfe/models/__init__.py +++ b/src/pytfe/models/__init__.py @@ -8,8 +8,8 @@ Agent, AgentListOptions, AgentPool, - AgentPoolAssignToWorkspacesOptions, AgentPoolAssignToProjectsOptions, + AgentPoolAssignToWorkspacesOptions, AgentPoolCreateOptions, AgentPoolListOptions, AgentPoolReadOptions,