From 0b7311d66c35d4749419836bab8ccbad8ad2cbbe Mon Sep 17 00:00:00 2001 From: "T.J Ariyawansa" Date: Fri, 1 May 2026 22:08:55 +0000 Subject: [PATCH 1/2] feat: add DatasetClient for Dataset Management API Implements DatasetClient mirroring GatewayClient pattern with single bedrock-agentcore-control boto3 client, __getattr__ dispatch for create/get/list/update/delete_dataset, and wait helpers. --- src/bedrock_agentcore/dataset/__init__.py | 3 + src/bedrock_agentcore/dataset/client.py | 155 ++++++++++++ tests/bedrock_agentcore/dataset/__init__.py | 0 .../bedrock_agentcore/dataset/test_client.py | 235 ++++++++++++++++++ 4 files changed, 393 insertions(+) create mode 100644 src/bedrock_agentcore/dataset/__init__.py create mode 100644 src/bedrock_agentcore/dataset/client.py create mode 100644 tests/bedrock_agentcore/dataset/__init__.py create mode 100644 tests/bedrock_agentcore/dataset/test_client.py diff --git a/src/bedrock_agentcore/dataset/__init__.py b/src/bedrock_agentcore/dataset/__init__.py new file mode 100644 index 00000000..ec99a6a6 --- /dev/null +++ b/src/bedrock_agentcore/dataset/__init__.py @@ -0,0 +1,3 @@ +from .client import DatasetClient + +__all__ = ["DatasetClient"] diff --git a/src/bedrock_agentcore/dataset/client.py b/src/bedrock_agentcore/dataset/client.py new file mode 100644 index 00000000..4edb8012 --- /dev/null +++ b/src/bedrock_agentcore/dataset/client.py @@ -0,0 +1,155 @@ +"""AgentCore Dataset SDK - Client for Dataset Management operations.""" + +import logging +from typing import Any, Dict, Optional + +import boto3 +from botocore.config import Config + +from .._utils.config import WaitConfig +from .._utils.polling import wait_until, wait_until_deleted +from .._utils.snake_case import accept_snake_case_kwargs, convert_kwargs +from .._utils.user_agent import build_user_agent_suffix + +logger = logging.getLogger(__name__) + +_DATASET_FAILED_STATUSES = {"FAILED", "UPDATE_UNSUCCESSFUL"} + + +class DatasetClient: + """Client for Bedrock AgentCore Dataset Management operations. + + Provides access to dataset CRUD operations. + Allowlisted boto3 methods can be called directly on this client. + Parameters accept both camelCase and snake_case (auto-converted). + + Example:: + + client = DatasetClient(region_name="us-west-2") + + # Pass-through to boto3 control plane client + dataset = client.create_dataset( + name="my-dataset", + roleArn="arn:aws:iam::123456789:role/dataset-role", + ) + """ + + _ALLOWED_CP_METHODS = { + "create_dataset", + "get_dataset", + "list_datasets", + "update_dataset", + "delete_dataset", + } + + def __init__( + self, + region_name: Optional[str] = None, + integration_source: Optional[str] = None, + boto3_session: Optional[boto3.Session] = None, + ): + """Initialize the Dataset client. + + Args: + region_name: AWS region name. If not provided, uses the session's region or "us-west-2". + integration_source: Optional integration source for user-agent telemetry. + boto3_session: Optional boto3 Session to use. If not provided, a default session + is created. Useful for named profiles or custom credentials. + """ + session = boto3_session if boto3_session else boto3.Session() + self.region_name = region_name or session.region_name or "us-west-2" + self.integration_source = integration_source + + user_agent_extra = build_user_agent_suffix(integration_source) + client_config = Config(user_agent_extra=user_agent_extra) + + self.cp_client = session.client( + "bedrock-agentcore-control", region_name=self.region_name, config=client_config + ) + + logger.info("Initialized DatasetClient for region: %s", self.cp_client.meta.region_name) + + # Pass-through + # ------------------------------------------------------------------------- + def __getattr__(self, name: str): + """Dynamically forward allowlisted method calls to the control plane boto3 client.""" + if name in self._ALLOWED_CP_METHODS and hasattr(self.cp_client, name): + method = getattr(self.cp_client, name) + logger.debug("Forwarding method '%s' to cp_client", name) + return accept_snake_case_kwargs(method) + + raise AttributeError( + f"'{self.__class__.__name__}' object has no attribute '{name}'. " + f"Method not found on cp_client. " + f"Available methods can be found in the boto3 documentation for " + f"'bedrock-agentcore-control' service." + ) + + # *_and_wait methods + # ------------------------------------------------------------------------- + def create_dataset_and_wait(self, wait_config: Optional[WaitConfig] = None, **kwargs) -> Dict[str, Any]: + """Create a dataset and wait for it to reach READY status. + + Args: + wait_config: Optional WaitConfig for polling behavior (default: max_wait=300, poll_interval=10). + **kwargs: Arguments forwarded to the create_dataset API. + + Returns: + Dataset details when READY. + + Raises: + RuntimeError: If the dataset reaches a failed state. + TimeoutError: If the dataset doesn't become READY within max_wait. + """ + response = self.cp_client.create_dataset(**convert_kwargs(kwargs)) + dataset_id = response["datasetId"] + return wait_until( + lambda: self.cp_client.get_dataset(datasetIdentifier=dataset_id), + "READY", + _DATASET_FAILED_STATUSES, + wait_config, + ) + + def update_dataset_and_wait(self, wait_config: Optional[WaitConfig] = None, **kwargs) -> Dict[str, Any]: + """Update a dataset and wait for it to reach READY status. + + Args: + wait_config: Optional WaitConfig for polling behavior (default: max_wait=300, poll_interval=10). + **kwargs: Arguments forwarded to the update_dataset API. + + Returns: + Dataset details when READY. + + Raises: + RuntimeError: If the dataset reaches a failed state. + TimeoutError: If the dataset doesn't become READY within max_wait. + """ + response = self.cp_client.update_dataset(**convert_kwargs(kwargs)) + dataset_id = response["datasetId"] + return wait_until( + lambda: self.cp_client.get_dataset(datasetIdentifier=dataset_id), + "READY", + _DATASET_FAILED_STATUSES, + wait_config, + ) + + def delete_dataset_and_wait( + self, + wait_config: Optional[WaitConfig] = None, + **kwargs, + ) -> None: + """Delete a dataset and wait for deletion to complete. + + Args: + wait_config: Optional WaitConfig for polling behavior. + **kwargs: Arguments forwarded to the delete_dataset API. + + Raises: + TimeoutError: If the dataset isn't deleted within max_wait. + """ + response = self.cp_client.delete_dataset(**convert_kwargs(kwargs)) + dataset_id = response["datasetId"] + wait_until_deleted( + lambda: self.cp_client.get_dataset(datasetIdentifier=dataset_id), + wait_config=wait_config, + ) diff --git a/tests/bedrock_agentcore/dataset/__init__.py b/tests/bedrock_agentcore/dataset/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/bedrock_agentcore/dataset/test_client.py b/tests/bedrock_agentcore/dataset/test_client.py new file mode 100644 index 00000000..a6b62b4c --- /dev/null +++ b/tests/bedrock_agentcore/dataset/test_client.py @@ -0,0 +1,235 @@ +"""Unit tests for DatasetClient - no external connections.""" + +from unittest.mock import MagicMock, patch + +import pytest +from botocore.exceptions import ClientError + +from bedrock_agentcore.dataset import DatasetClient + + +def test_client_initialization(): + """Test client initialization creates a single bedrock-agentcore-control client.""" + with patch("boto3.Session") as mock_session_cls: + mock_session = MagicMock() + mock_session.region_name = "us-west-2" + mock_session_cls.return_value = mock_session + + client = DatasetClient(region_name="us-west-2") + + assert client.region_name == "us-west-2" + assert client.integration_source is None + # Only one client: bedrock-agentcore-control + assert mock_session.client.call_count == 1 + + # Verify config was passed + for call in mock_session.client.call_args_list: + assert "config" in call.kwargs + + # Verify the correct service was used + service_name = mock_session.client.call_args_list[0][0][0] + assert service_name == "bedrock-agentcore-control" + + +def test_client_initialization_with_integration_source(): + """Test client initialization with integration_source sets user-agent telemetry.""" + with patch("boto3.Session") as mock_session_cls: + mock_session = MagicMock() + mock_session.region_name = "us-west-2" + mock_session_cls.return_value = mock_session + + client = DatasetClient(region_name="us-west-2", integration_source="langchain") + + assert client.region_name == "us-west-2" + assert client.integration_source == "langchain" + assert mock_session.client.call_count == 1 + + +def test_client_initialization_with_boto3_session(): + """Test client initialization with a custom boto3 session.""" + mock_session = MagicMock() + mock_session.region_name = "eu-west-1" + + client = DatasetClient(boto3_session=mock_session) + + assert client.region_name == "eu-west-1" + assert mock_session.client.call_count == 1 + + # Verify the correct service was used + call_args = [call[0][0] for call in mock_session.client.call_args_list] + assert "bedrock-agentcore-control" in call_args + + +def test_client_initialization_region_fallback_to_default(): + """Test region falls back to us-west-2 when no region is provided.""" + with patch("boto3.Session") as mock_session_cls: + mock_session = MagicMock() + mock_session.region_name = None + mock_session_cls.return_value = mock_session + + client = DatasetClient() + + assert client.region_name == "us-west-2" + + +def test_getattr_allowed_methods(): + """Verify each allowed method dispatches correctly via __getattr__.""" + mock_session = MagicMock() + mock_boto_client = MagicMock() + mock_session.client.return_value = mock_boto_client + + client = DatasetClient(region_name="us-east-1", boto3_session=mock_session) + + allowed_methods = [ + "create_dataset", + "get_dataset", + "list_datasets", + "update_dataset", + "delete_dataset", + ] + + for method_name in allowed_methods: + method = getattr(client, method_name) + assert callable(method), f"{method_name} should be callable" + + +def test_getattr_unknown_method_raises(): + """AttributeError is raised for unknown/disallowed methods.""" + mock_session = MagicMock() + mock_session.client.return_value = MagicMock() + + client = DatasetClient(region_name="us-east-1", boto3_session=mock_session) + + with pytest.raises(AttributeError, match="has no attribute 'create_evaluator'"): + _ = client.create_evaluator + + +def test_create_dataset(): + """Test create_dataset delegates to cp_client.""" + mock_session = MagicMock() + mock_boto_client = MagicMock() + mock_session.client.return_value = mock_boto_client + mock_boto_client.create_dataset.return_value = {"datasetId": "ds-123", "status": "CREATING"} + + client = DatasetClient(region_name="us-west-2", boto3_session=mock_session) + result = client.create_dataset(name="my-dataset") + + mock_boto_client.create_dataset.assert_called_once_with(name="my-dataset") + assert result["datasetId"] == "ds-123" + + +def test_list_datasets(): + """Test list_datasets delegates to cp_client.""" + mock_session = MagicMock() + mock_boto_client = MagicMock() + mock_session.client.return_value = mock_boto_client + mock_boto_client.list_datasets.return_value = { + "items": [{"datasetId": "ds-1"}, {"datasetId": "ds-2"}] + } + + client = DatasetClient(region_name="us-west-2", boto3_session=mock_session) + result = client.list_datasets() + + mock_boto_client.list_datasets.assert_called_once_with() + assert len(result["items"]) == 2 + + +def test_delete_dataset(): + """Test delete_dataset delegates to cp_client.""" + mock_session = MagicMock() + mock_boto_client = MagicMock() + mock_session.client.return_value = mock_boto_client + mock_boto_client.delete_dataset.return_value = {"datasetId": "ds-123"} + + client = DatasetClient(region_name="us-west-2", boto3_session=mock_session) + result = client.delete_dataset(datasetIdentifier="ds-123") + + mock_boto_client.delete_dataset.assert_called_once_with(datasetIdentifier="ds-123") + assert result["datasetId"] == "ds-123" + + +def test_delete_dataset_and_wait(): + """Test delete_dataset_and_wait polls until the dataset is deleted.""" + from botocore.exceptions import ClientError + + mock_session = MagicMock() + mock_boto_client = MagicMock() + mock_session.client.return_value = mock_boto_client + + # delete_dataset returns the datasetId + mock_boto_client.delete_dataset.return_value = {"datasetId": "ds-123"} + + # get_dataset raises ResourceNotFoundException after deletion + not_found_error = ClientError( + {"Error": {"Code": "ResourceNotFoundException", "Message": "Not found"}}, + "GetDataset", + ) + mock_boto_client.get_dataset.side_effect = not_found_error + + client = DatasetClient(region_name="us-west-2", boto3_session=mock_session) + # Should not raise + client.delete_dataset_and_wait(datasetIdentifier="ds-123") + + mock_boto_client.delete_dataset.assert_called_once() + mock_boto_client.get_dataset.assert_called_once_with(datasetIdentifier="ds-123") + + +def test_create_dataset_and_wait(): + """Test create_dataset_and_wait polls until READY status.""" + mock_session = MagicMock() + mock_boto_client = MagicMock() + mock_session.client.return_value = mock_boto_client + + mock_boto_client.create_dataset.return_value = {"datasetId": "ds-456", "status": "CREATING"} + mock_boto_client.get_dataset.return_value = {"datasetId": "ds-456", "status": "READY"} + + client = DatasetClient(region_name="us-west-2", boto3_session=mock_session) + result = client.create_dataset_and_wait(name="my-dataset") + + assert result["status"] == "READY" + assert result["datasetId"] == "ds-456" + mock_boto_client.create_dataset.assert_called_once_with(name="my-dataset") + mock_boto_client.get_dataset.assert_called_once_with(datasetIdentifier="ds-456") + + +def test_update_dataset_and_wait(): + """Test update_dataset_and_wait polls until READY status.""" + mock_session = MagicMock() + mock_boto_client = MagicMock() + mock_session.client.return_value = mock_boto_client + + mock_boto_client.update_dataset.return_value = {"datasetId": "ds-789", "status": "UPDATING"} + mock_boto_client.get_dataset.return_value = {"datasetId": "ds-789", "status": "READY"} + + client = DatasetClient(region_name="us-west-2", boto3_session=mock_session) + result = client.update_dataset_and_wait(datasetIdentifier="ds-789", name="updated-name") + + assert result["status"] == "READY" + mock_boto_client.update_dataset.assert_called_once_with(datasetIdentifier="ds-789", name="updated-name") + mock_boto_client.get_dataset.assert_called_once_with(datasetIdentifier="ds-789") + + +def test_create_dataset_and_wait_failed_status(): + """Test create_dataset_and_wait raises RuntimeError on FAILED status.""" + mock_session = MagicMock() + mock_boto_client = MagicMock() + mock_session.client.return_value = mock_boto_client + + mock_boto_client.create_dataset.return_value = {"datasetId": "ds-bad", "status": "CREATING"} + mock_boto_client.get_dataset.return_value = {"datasetId": "ds-bad", "status": "FAILED", "statusReasons": "Bad config"} + + client = DatasetClient(region_name="us-west-2", boto3_session=mock_session) + + with pytest.raises(RuntimeError, match="FAILED"): + client.create_dataset_and_wait(name="bad-dataset") + + +def test_getattr_error_message_contains_service(): + """AttributeError message mentions the service name.""" + mock_session = MagicMock() + mock_session.client.return_value = MagicMock() + + client = DatasetClient(region_name="us-east-1", boto3_session=mock_session) + + with pytest.raises(AttributeError, match="bedrock-agentcore-control"): + _ = client.unknown_operation From 325dd825598b8d656140559fca840c5ae1159164 Mon Sep 17 00:00:00 2001 From: "T.J Ariyawansa" Date: Fri, 1 May 2026 23:40:34 +0000 Subject: [PATCH 2/2] fix: address critical and high issues in DatasetClient - Export DatasetClient from top-level bedrock_agentcore package - Unwrap response["dataset"] in polling lambdas for create/update_and_wait - Add clientToken (uuid4) idempotency to create_dataset_and_wait - Guard delete_dataset_and_wait with response.get("datasetId") + ValueError - Add get_paginator and get_waiter to _ALLOWED_CP_METHODS - Update tests to 100% coverage (24 tests) covering all new behaviours Co-Authored-By: Claude Sonnet 4.6 (1M context) --- src/bedrock_agentcore/__init__.py | 2 + src/bedrock_agentcore/dataset/client.py | 15 +- .../bedrock_agentcore/dataset/test_client.py | 313 +++++++++++++----- 3 files changed, 239 insertions(+), 91 deletions(-) diff --git a/src/bedrock_agentcore/__init__.py b/src/bedrock_agentcore/__init__.py index a77472f5..d46aaf02 100644 --- a/src/bedrock_agentcore/__init__.py +++ b/src/bedrock_agentcore/__init__.py @@ -1,5 +1,6 @@ """BedrockAgentCore Runtime SDK - A Python SDK for building and deploying AI agents.""" +from .dataset import DatasetClient from .runtime import BedrockAgentCoreApp, BedrockAgentCoreContext, RequestContext from .runtime.models import PingStatus @@ -8,4 +9,5 @@ "RequestContext", "BedrockAgentCoreContext", "PingStatus", + "DatasetClient", ] diff --git a/src/bedrock_agentcore/dataset/client.py b/src/bedrock_agentcore/dataset/client.py index 4edb8012..be7ccab8 100644 --- a/src/bedrock_agentcore/dataset/client.py +++ b/src/bedrock_agentcore/dataset/client.py @@ -1,6 +1,7 @@ """AgentCore Dataset SDK - Client for Dataset Management operations.""" import logging +import uuid from typing import Any, Dict, Optional import boto3 @@ -40,6 +41,8 @@ class DatasetClient: "list_datasets", "update_dataset", "delete_dataset", + "get_paginator", + "get_waiter", } def __init__( @@ -101,10 +104,12 @@ def create_dataset_and_wait(self, wait_config: Optional[WaitConfig] = None, **kw RuntimeError: If the dataset reaches a failed state. TimeoutError: If the dataset doesn't become READY within max_wait. """ - response = self.cp_client.create_dataset(**convert_kwargs(kwargs)) + params = convert_kwargs(kwargs) + params.setdefault("clientToken", str(uuid.uuid4())) + response = self.cp_client.create_dataset(**params) dataset_id = response["datasetId"] return wait_until( - lambda: self.cp_client.get_dataset(datasetIdentifier=dataset_id), + lambda: self.cp_client.get_dataset(datasetIdentifier=dataset_id)["dataset"], "READY", _DATASET_FAILED_STATUSES, wait_config, @@ -127,7 +132,7 @@ def update_dataset_and_wait(self, wait_config: Optional[WaitConfig] = None, **kw response = self.cp_client.update_dataset(**convert_kwargs(kwargs)) dataset_id = response["datasetId"] return wait_until( - lambda: self.cp_client.get_dataset(datasetIdentifier=dataset_id), + lambda: self.cp_client.get_dataset(datasetIdentifier=dataset_id)["dataset"], "READY", _DATASET_FAILED_STATUSES, wait_config, @@ -148,7 +153,9 @@ def delete_dataset_and_wait( TimeoutError: If the dataset isn't deleted within max_wait. """ response = self.cp_client.delete_dataset(**convert_kwargs(kwargs)) - dataset_id = response["datasetId"] + dataset_id = response.get("datasetId") + if not dataset_id: + raise ValueError("delete_dataset response did not include a 'datasetId'; cannot poll for deletion.") wait_until_deleted( lambda: self.cp_client.get_dataset(datasetIdentifier=dataset_id), wait_config=wait_config, diff --git a/tests/bedrock_agentcore/dataset/test_client.py b/tests/bedrock_agentcore/dataset/test_client.py index a6b62b4c..0020f012 100644 --- a/tests/bedrock_agentcore/dataset/test_client.py +++ b/tests/bedrock_agentcore/dataset/test_client.py @@ -1,6 +1,6 @@ """Unit tests for DatasetClient - no external connections.""" -from unittest.mock import MagicMock, patch +from unittest.mock import MagicMock, call, patch import pytest from botocore.exceptions import ClientError @@ -8,6 +8,31 @@ from bedrock_agentcore.dataset import DatasetClient +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +def _make_client(region="us-west-2", **kwargs): + """Return a DatasetClient backed by a single MagicMock boto3 client.""" + mock_session = MagicMock() + mock_boto_client = MagicMock() + mock_session.region_name = region + mock_session.client.return_value = mock_boto_client + client = DatasetClient(region_name=region, boto3_session=mock_session, **kwargs) + return client, mock_boto_client + + +def _not_found_error(): + return ClientError( + {"Error": {"Code": "ResourceNotFoundException", "Message": "Not found"}}, + "GetDataset", + ) + + +# --------------------------------------------------------------------------- +# Initialisation +# --------------------------------------------------------------------------- + def test_client_initialization(): """Test client initialization creates a single bedrock-agentcore-control client.""" with patch("boto3.Session") as mock_session_cls: @@ -23,8 +48,8 @@ def test_client_initialization(): assert mock_session.client.call_count == 1 # Verify config was passed - for call in mock_session.client.call_args_list: - assert "config" in call.kwargs + for call_item in mock_session.client.call_args_list: + assert "config" in call_item.kwargs # Verify the correct service was used service_name = mock_session.client.call_args_list[0][0][0] @@ -56,7 +81,7 @@ def test_client_initialization_with_boto3_session(): assert mock_session.client.call_count == 1 # Verify the correct service was used - call_args = [call[0][0] for call in mock_session.client.call_args_list] + call_args = [c[0][0] for c in mock_session.client.call_args_list] assert "bedrock-agentcore-control" in call_args @@ -72,13 +97,23 @@ def test_client_initialization_region_fallback_to_default(): assert client.region_name == "us-west-2" +# --------------------------------------------------------------------------- +# Top-level export +# --------------------------------------------------------------------------- + +def test_dataset_client_exported_from_top_level(): + """DatasetClient must be importable directly from bedrock_agentcore.""" + from bedrock_agentcore import DatasetClient as DC # noqa: F401 + assert DC is DatasetClient + + +# --------------------------------------------------------------------------- +# __getattr__ dispatch +# --------------------------------------------------------------------------- + def test_getattr_allowed_methods(): """Verify each allowed method dispatches correctly via __getattr__.""" - mock_session = MagicMock() - mock_boto_client = MagicMock() - mock_session.client.return_value = mock_boto_client - - client = DatasetClient(region_name="us-east-1", boto3_session=mock_session) + client, _ = _make_client() allowed_methods = [ "create_dataset", @@ -86,6 +121,8 @@ def test_getattr_allowed_methods(): "list_datasets", "update_dataset", "delete_dataset", + "get_paginator", + "get_waiter", ] for method_name in allowed_methods: @@ -95,141 +132,243 @@ def test_getattr_allowed_methods(): def test_getattr_unknown_method_raises(): """AttributeError is raised for unknown/disallowed methods.""" - mock_session = MagicMock() - mock_session.client.return_value = MagicMock() - - client = DatasetClient(region_name="us-east-1", boto3_session=mock_session) + client, _ = _make_client() with pytest.raises(AttributeError, match="has no attribute 'create_evaluator'"): _ = client.create_evaluator +def test_getattr_error_message_contains_service(): + """AttributeError message mentions the service name.""" + client, _ = _make_client() + + with pytest.raises(AttributeError, match="bedrock-agentcore-control"): + _ = client.unknown_operation + + +# --------------------------------------------------------------------------- +# Pass-through CRUD +# --------------------------------------------------------------------------- + def test_create_dataset(): """Test create_dataset delegates to cp_client.""" - mock_session = MagicMock() - mock_boto_client = MagicMock() - mock_session.client.return_value = mock_boto_client - mock_boto_client.create_dataset.return_value = {"datasetId": "ds-123", "status": "CREATING"} + client, mock_boto = _make_client() + mock_boto.create_dataset.return_value = {"datasetId": "ds-123", "status": "CREATING"} - client = DatasetClient(region_name="us-west-2", boto3_session=mock_session) result = client.create_dataset(name="my-dataset") - mock_boto_client.create_dataset.assert_called_once_with(name="my-dataset") + mock_boto.create_dataset.assert_called_once_with(name="my-dataset") assert result["datasetId"] == "ds-123" def test_list_datasets(): """Test list_datasets delegates to cp_client.""" - mock_session = MagicMock() - mock_boto_client = MagicMock() - mock_session.client.return_value = mock_boto_client - mock_boto_client.list_datasets.return_value = { + client, mock_boto = _make_client() + mock_boto.list_datasets.return_value = { "items": [{"datasetId": "ds-1"}, {"datasetId": "ds-2"}] } - client = DatasetClient(region_name="us-west-2", boto3_session=mock_session) result = client.list_datasets() - mock_boto_client.list_datasets.assert_called_once_with() + mock_boto.list_datasets.assert_called_once_with() assert len(result["items"]) == 2 def test_delete_dataset(): """Test delete_dataset delegates to cp_client.""" - mock_session = MagicMock() - mock_boto_client = MagicMock() - mock_session.client.return_value = mock_boto_client - mock_boto_client.delete_dataset.return_value = {"datasetId": "ds-123"} + client, mock_boto = _make_client() + mock_boto.delete_dataset.return_value = {"datasetId": "ds-123"} - client = DatasetClient(region_name="us-west-2", boto3_session=mock_session) result = client.delete_dataset(datasetIdentifier="ds-123") - mock_boto_client.delete_dataset.assert_called_once_with(datasetIdentifier="ds-123") + mock_boto.delete_dataset.assert_called_once_with(datasetIdentifier="ds-123") assert result["datasetId"] == "ds-123" -def test_delete_dataset_and_wait(): - """Test delete_dataset_and_wait polls until the dataset is deleted.""" - from botocore.exceptions import ClientError +# --------------------------------------------------------------------------- +# create_dataset_and_wait +# --------------------------------------------------------------------------- - mock_session = MagicMock() - mock_boto_client = MagicMock() - mock_session.client.return_value = mock_boto_client +def test_create_dataset_and_wait(): + """Test create_dataset_and_wait polls until READY status.""" + client, mock_boto = _make_client() - # delete_dataset returns the datasetId - mock_boto_client.delete_dataset.return_value = {"datasetId": "ds-123"} + mock_boto.create_dataset.return_value = {"datasetId": "ds-456", "status": "CREATING"} + mock_boto.get_dataset.return_value = {"dataset": {"datasetId": "ds-456", "status": "READY"}} - # get_dataset raises ResourceNotFoundException after deletion - not_found_error = ClientError( - {"Error": {"Code": "ResourceNotFoundException", "Message": "Not found"}}, - "GetDataset", - ) - mock_boto_client.get_dataset.side_effect = not_found_error + result = client.create_dataset_and_wait(name="my-dataset") - client = DatasetClient(region_name="us-west-2", boto3_session=mock_session) - # Should not raise - client.delete_dataset_and_wait(datasetIdentifier="ds-123") + assert result["status"] == "READY" + assert result["datasetId"] == "ds-456" + mock_boto.create_dataset.assert_called_once() + mock_boto.get_dataset.assert_called_once_with(datasetIdentifier="ds-456") - mock_boto_client.delete_dataset.assert_called_once() - mock_boto_client.get_dataset.assert_called_once_with(datasetIdentifier="ds-123") +def test_create_dataset_and_wait_adds_client_token(): + """create_dataset_and_wait injects a clientToken for idempotency.""" + client, mock_boto = _make_client() -def test_create_dataset_and_wait(): - """Test create_dataset_and_wait polls until READY status.""" - mock_session = MagicMock() - mock_boto_client = MagicMock() - mock_session.client.return_value = mock_boto_client + mock_boto.create_dataset.return_value = {"datasetId": "ds-tok", "status": "CREATING"} + mock_boto.get_dataset.return_value = {"dataset": {"datasetId": "ds-tok", "status": "READY"}} - mock_boto_client.create_dataset.return_value = {"datasetId": "ds-456", "status": "CREATING"} - mock_boto_client.get_dataset.return_value = {"datasetId": "ds-456", "status": "READY"} + client.create_dataset_and_wait(name="tok-dataset") - client = DatasetClient(region_name="us-west-2", boto3_session=mock_session) - result = client.create_dataset_and_wait(name="my-dataset") + create_call_kwargs = mock_boto.create_dataset.call_args[1] + assert "clientToken" in create_call_kwargs + # Should be a non-empty string (UUID) + assert len(create_call_kwargs["clientToken"]) > 0 + + +def test_create_dataset_and_wait_respects_explicit_client_token(): + """Caller-supplied clientToken must not be overwritten.""" + client, mock_boto = _make_client() + + mock_boto.create_dataset.return_value = {"datasetId": "ds-ct", "status": "CREATING"} + mock_boto.get_dataset.return_value = {"dataset": {"datasetId": "ds-ct", "status": "READY"}} + + client.create_dataset_and_wait(name="ct-dataset", clientToken="my-token-123") + + create_call_kwargs = mock_boto.create_dataset.call_args[1] + assert create_call_kwargs["clientToken"] == "my-token-123" + + +def test_create_dataset_and_wait_failed_status(): + """Test create_dataset_and_wait raises RuntimeError on FAILED status.""" + client, mock_boto = _make_client() + + mock_boto.create_dataset.return_value = {"datasetId": "ds-bad", "status": "CREATING"} + mock_boto.get_dataset.return_value = { + "dataset": {"datasetId": "ds-bad", "status": "FAILED", "statusReasons": "Bad config"} + } + + with pytest.raises(RuntimeError, match="FAILED"): + client.create_dataset_and_wait(name="bad-dataset") + + +def test_create_dataset_and_wait_update_unsuccessful(): + """create_dataset_and_wait raises RuntimeError on UPDATE_UNSUCCESSFUL status.""" + client, mock_boto = _make_client() + + mock_boto.create_dataset.return_value = {"datasetId": "ds-upd", "status": "CREATING"} + mock_boto.get_dataset.return_value = { + "dataset": {"datasetId": "ds-upd", "status": "UPDATE_UNSUCCESSFUL"} + } + + with pytest.raises(RuntimeError, match="UPDATE_UNSUCCESSFUL"): + client.create_dataset_and_wait(name="upd-dataset") + + +def test_create_dataset_and_wait_polls_nested_dataset_key(): + """The poll lambda must unwrap response['dataset'] before reading status.""" + client, mock_boto = _make_client() + + mock_boto.create_dataset.return_value = {"datasetId": "ds-nested", "status": "CREATING"} + # First call returns CREATING, second returns READY — both wrapped under 'dataset' + mock_boto.get_dataset.side_effect = [ + {"dataset": {"datasetId": "ds-nested", "status": "CREATING"}}, + {"dataset": {"datasetId": "ds-nested", "status": "READY"}}, + ] + + result = client.create_dataset_and_wait(name="nested-dataset") assert result["status"] == "READY" - assert result["datasetId"] == "ds-456" - mock_boto_client.create_dataset.assert_called_once_with(name="my-dataset") - mock_boto_client.get_dataset.assert_called_once_with(datasetIdentifier="ds-456") + assert mock_boto.get_dataset.call_count == 2 + +# --------------------------------------------------------------------------- +# update_dataset_and_wait +# --------------------------------------------------------------------------- def test_update_dataset_and_wait(): """Test update_dataset_and_wait polls until READY status.""" - mock_session = MagicMock() - mock_boto_client = MagicMock() - mock_session.client.return_value = mock_boto_client + client, mock_boto = _make_client() - mock_boto_client.update_dataset.return_value = {"datasetId": "ds-789", "status": "UPDATING"} - mock_boto_client.get_dataset.return_value = {"datasetId": "ds-789", "status": "READY"} + mock_boto.update_dataset.return_value = {"datasetId": "ds-789", "status": "UPDATING"} + mock_boto.get_dataset.return_value = {"dataset": {"datasetId": "ds-789", "status": "READY"}} - client = DatasetClient(region_name="us-west-2", boto3_session=mock_session) result = client.update_dataset_and_wait(datasetIdentifier="ds-789", name="updated-name") assert result["status"] == "READY" - mock_boto_client.update_dataset.assert_called_once_with(datasetIdentifier="ds-789", name="updated-name") - mock_boto_client.get_dataset.assert_called_once_with(datasetIdentifier="ds-789") + mock_boto.update_dataset.assert_called_once_with(datasetIdentifier="ds-789", name="updated-name") + mock_boto.get_dataset.assert_called_once_with(datasetIdentifier="ds-789") -def test_create_dataset_and_wait_failed_status(): - """Test create_dataset_and_wait raises RuntimeError on FAILED status.""" - mock_session = MagicMock() - mock_boto_client = MagicMock() - mock_session.client.return_value = mock_boto_client +def test_update_dataset_and_wait_unwraps_nested_key(): + """update_dataset_and_wait must unwrap response['dataset'] for status check.""" + client, mock_boto = _make_client() - mock_boto_client.create_dataset.return_value = {"datasetId": "ds-bad", "status": "CREATING"} - mock_boto_client.get_dataset.return_value = {"datasetId": "ds-bad", "status": "FAILED", "statusReasons": "Bad config"} + mock_boto.update_dataset.return_value = {"datasetId": "ds-upd2", "status": "UPDATING"} + mock_boto.get_dataset.side_effect = [ + {"dataset": {"datasetId": "ds-upd2", "status": "UPDATING"}}, + {"dataset": {"datasetId": "ds-upd2", "status": "READY"}}, + ] - client = DatasetClient(region_name="us-west-2", boto3_session=mock_session) + result = client.update_dataset_and_wait(datasetIdentifier="ds-upd2", name="v2") - with pytest.raises(RuntimeError, match="FAILED"): - client.create_dataset_and_wait(name="bad-dataset") + assert result["status"] == "READY" + assert mock_boto.get_dataset.call_count == 2 -def test_getattr_error_message_contains_service(): - """AttributeError message mentions the service name.""" - mock_session = MagicMock() - mock_session.client.return_value = MagicMock() +# --------------------------------------------------------------------------- +# delete_dataset_and_wait +# --------------------------------------------------------------------------- - client = DatasetClient(region_name="us-east-1", boto3_session=mock_session) +def test_delete_dataset_and_wait(): + """Test delete_dataset_and_wait polls until the dataset is deleted.""" + client, mock_boto = _make_client() - with pytest.raises(AttributeError, match="bedrock-agentcore-control"): - _ = client.unknown_operation + mock_boto.delete_dataset.return_value = {"datasetId": "ds-123"} + mock_boto.get_dataset.side_effect = _not_found_error() + + client.delete_dataset_and_wait(datasetIdentifier="ds-123") + + mock_boto.delete_dataset.assert_called_once() + mock_boto.get_dataset.assert_called_once_with(datasetIdentifier="ds-123") + + +def test_delete_dataset_and_wait_missing_dataset_id_raises(): + """delete_dataset_and_wait raises ValueError when response has no datasetId.""" + client, mock_boto = _make_client() + + mock_boto.delete_dataset.return_value = {} # missing datasetId + + with pytest.raises(ValueError, match="datasetId"): + client.delete_dataset_and_wait(datasetIdentifier="ds-missing") + + +def test_delete_dataset_and_wait_non_404_error_propagates(): + """Non-ResourceNotFoundException errors bubble up during polling.""" + client, mock_boto = _make_client() + + mock_boto.delete_dataset.return_value = {"datasetId": "ds-err"} + mock_boto.get_dataset.side_effect = ClientError( + {"Error": {"Code": "InternalServerError", "Message": "Oops"}}, + "GetDataset", + ) + + with pytest.raises(ClientError, match="InternalServerError"): + client.delete_dataset_and_wait(datasetIdentifier="ds-err") + + +# --------------------------------------------------------------------------- +# Pagination helpers +# --------------------------------------------------------------------------- + +def test_get_paginator_is_accessible(): + """get_paginator must be reachable via __getattr__.""" + client, mock_boto = _make_client() + mock_boto.get_paginator.return_value = MagicMock() + + paginator = client.get_paginator("list_datasets") + + mock_boto.get_paginator.assert_called_once_with("list_datasets") + + +def test_get_waiter_is_accessible(): + """get_waiter must be reachable via __getattr__.""" + client, mock_boto = _make_client() + mock_boto.get_waiter.return_value = MagicMock() + + waiter = client.get_waiter("dataset_ready") + + mock_boto.get_waiter.assert_called_once_with("dataset_ready")