add proxy tests for Cognito-IDP

whummer · whummer · commit 1a9fa00b7b3b · 2025-11-15T10:50:29.000-05:00
diff --git a/aws-proxy/AGENTS.md b/aws-proxy/AGENTS.md
@@ -1,13 +1,31 @@
-# AI Agent Instructions for AWS Proxy Extension
+# AI Agent Instructions for LocalStack AWS Proxy Extension
 
 This repo is a LocalStack extension (plugin) that enables a "proxy mode" - proxying requests for certain AWS services (e.g., S3) to the upstream real AWS cloud, while handling the remaining services locally.
 
+You are an AI agent tasked with adding additional functionality or test coverage to this repo.
+
+## General Instructions
+
+* You can assume that the LocalStack container is running locally, with the proxy extension installed and enabled.
+* You can assume that test AWS credentials are configured in the shell environment where the AI agent is running.
+* Do *not* touch any files outside the working directory (basedir of this file)!
+* You can create new files (no need to prompt for confirmation)
+* You can make modifications to files (no need to prompt for confirmation)
+* You can delete existing files if needed, but only after user confirmation
+* You can call different `make` targets (e.g., `make test`) in this repo (no need to prompt for confirmation)
+* For each new file created or existing file modified, add a header comment to the file, something like `# Note/disclosure: This file has been (partially or fully) generated by an AI agent.`
+* The proxy tests are executed against real AWS and may incur some costs, so rather than executing the entire test suite or entire modules, focus the testing on individual test functions within a module only.
+* To format/lint the codebase you can run `make format` and `make lint`
+
 ## Testing
 
-The proxy functionality is covered by integration tests in the `tests/` folder, one file for each different service.
+The proxy functionality is covered by integration tests in the `tests/` folder, one file for each different AWS service (e.g., SQS, S3, etc).
 
 To add a test, follow the pattern in the existing tests.
 It usually involves creating two boto3 clients, one for the LocalStack connection, and one for the real upstream AWS cloud.
 We then run API requests with both clients and assert that the results are identical, thereby ensuring that the proxy functionality is working properly.
 
-You can assume that test AWS credentials are configured in the shell environment where the AI agent is running.
+To run a single test via `pytest` (say, `test_my_logic` in `test_s3.py`), use the following command:
+```
+TEST_PATH=tests/test_s3.py::test_my_logic make test
+```
diff --git a/aws-proxy/pyproject.toml b/aws-proxy/pyproject.toml
@@ -28,6 +28,7 @@ test=[
     "pytest-httpserver",
     "ruff",
     # build tools
+    "build",
     "setuptools",
     "setuptools_scm",
     "wheel"
diff --git a/aws-proxy/tests/proxy/test_cognito_idp.py b/aws-proxy/tests/proxy/test_cognito_idp.py
@@ -0,0 +1,164 @@
+# Note/disclosure: This file has been (partially or fully) generated by an AI agent.
+
+import boto3
+import pytest
+from botocore.exceptions import ClientError
+from localstack.aws.connect import connect_to
+from localstack.utils.sync import retry
+
+from aws_proxy.shared.models import ProxyConfig
+
+
+def _compare_user_pool_lists(user_pools_aws: list, user_pools_proxied: list):
+    # Helper to compare user pool lists, ignoring some dynamic attributes
+    # We'll need to refine this based on actual Cognito-IDP response structure
+    def normalize_user_pool(user_pool):
+        normalized = user_pool.copy()
+        normalized.pop("LastModifiedDate", None)
+        normalized.pop("CreationDate", None)
+        return normalized
+
+    normalized_aws = sorted(
+        [normalize_user_pool(up) for up in user_pools_aws], key=lambda x: x["Id"]
+    )
+    normalized_proxied = sorted(
+        [normalize_user_pool(up) for up in user_pools_proxied], key=lambda x: x["Id"]
+    )
+
+    assert normalized_proxied == normalized_aws
+
+
+def test_cognito_idp_user_pool_operations(start_aws_proxy):
+    # Start proxy for cognito-idp service
+    config = ProxyConfig(services={"cognito-idp": {"resources": ".*"}})
+    start_aws_proxy(config)
+
+    # Create clients
+    cognito_client = connect_to().cognito_idp
+    cognito_client_aws = boto3.client("cognito-idp")
+
+    # List user pools initially to ensure clean state and compare
+    user_pools_proxied_initial = cognito_client.list_user_pools(MaxResults=10)[
+        "UserPools"
+    ]
+    user_pools_aws_initial = cognito_client_aws.list_user_pools(MaxResults=10)[
+        "UserPools"
+    ]
+    _compare_user_pool_lists(user_pools_aws_initial, user_pools_proxied_initial)
+
+    # Create a user pool
+    user_pool_name = "test-user-pool-proxy"
+    create_response_proxied = cognito_client.create_user_pool(PoolName=user_pool_name)
+    user_pool_id_proxied = create_response_proxied["UserPool"]["Id"]
+
+    # List user pools and compare
+    user_pools_proxied_after_create = cognito_client.list_user_pools(MaxResults=10)[
+        "UserPools"
+    ]
+    user_pools_aws_after_create = cognito_client_aws.list_user_pools(MaxResults=10)[
+        "UserPools"
+    ]
+    _compare_user_pool_lists(
+        user_pools_aws_after_create, user_pools_proxied_after_create
+    )
+
+    # Describe user pool and compare
+    describe_proxied = cognito_client.describe_user_pool(
+        UserPoolId=user_pool_id_proxied
+    )["UserPool"]
+    describe_aws = cognito_client_aws.describe_user_pool(
+        UserPoolId=user_pool_id_proxied
+    )["UserPool"]
+    # Normalize for comparison (remove dynamic fields)
+    describe_proxied.pop("LastModifiedDate", None)
+    describe_proxied.pop("CreationDate", None)
+    describe_aws.pop("LastModifiedDate", None)
+    describe_aws.pop("CreationDate", None)
+    assert describe_proxied == describe_aws
+
+    # Delete the user pool
+    cognito_client.delete_user_pool(UserPoolId=user_pool_id_proxied)
+
+    def _assert_user_pool_deleted():
+        with pytest.raises(ClientError) as aws_exc:
+            cognito_client_aws.describe_user_pool(UserPoolId=user_pool_id_proxied)
+        with pytest.raises(ClientError) as proxied_exc:
+            cognito_client.describe_user_pool(UserPoolId=user_pool_id_proxied)
+        assert aws_exc.value.response["Error"]["Code"] == "ResourceNotFoundException"
+        assert (
+            proxied_exc.value.response["Error"]["Code"] == "ResourceNotFoundException"
+        )
+
+    retry(_assert_user_pool_deleted, retries=5, sleep=5)
+
+    # Verify that the user pool is no longer in the list
+    user_pools_proxied_final = cognito_client.list_user_pools(MaxResults=10)[
+        "UserPools"
+    ]
+    user_pools_aws_final = cognito_client_aws.list_user_pools(MaxResults=10)[
+        "UserPools"
+    ]
+    _compare_user_pool_lists(user_pools_aws_final, user_pools_proxied_final)
+
+
+def test_cognito_idp_user_operations(start_aws_proxy):
+    # Start proxy for cognito-idp service
+    config = ProxyConfig(services={"cognito-idp": {"resources": ".*"}})
+    start_aws_proxy(config)
+
+    # Create clients
+    cognito_client = connect_to().cognito_idp
+    cognito_client_aws = boto3.client("cognito-idp")
+
+    # Create a user pool
+    user_pool_name = "test-user-pool-users-proxy"
+    user_pool = cognito_client.create_user_pool(PoolName=user_pool_name)
+    user_pool_id = user_pool["UserPool"]["Id"]
+
+    try:
+        # Create a user pool client
+        user_pool_client = cognito_client.create_user_pool_client(
+            UserPoolId=user_pool_id,
+            ClientName="test-client",
+            ExplicitAuthFlows=["ADMIN_NO_SRP_AUTH", "USER_PASSWORD_AUTH"],
+        )
+        client_id = user_pool_client["UserPoolClient"]["ClientId"]
+
+        # Sign up a user
+        username = "testuser"
+        password = "Password123!"
+        cognito_client.sign_up(
+            ClientId=client_id,
+            Username=username,
+            Password=password,
+            UserAttributes=[{"Name": "email", "Value": "test@example.com"}],
+        )
+
+        # Confirm the user
+        cognito_client.admin_confirm_sign_up(UserPoolId=user_pool_id, Username=username)
+
+        # Initiate auth
+        auth_response_proxied = cognito_client.initiate_auth(
+            ClientId=client_id,
+            AuthFlow="USER_PASSWORD_AUTH",
+            AuthParameters={"USERNAME": username, "PASSWORD": password},
+        )
+        auth_response_aws = cognito_client_aws.initiate_auth(
+            ClientId=client_id,
+            AuthFlow="USER_PASSWORD_AUTH",
+            AuthParameters={"USERNAME": username, "PASSWORD": password},
+        )
+        assert auth_response_proxied["AuthenticationResult"]["AccessToken"]
+        assert auth_response_aws["AuthenticationResult"]["AccessToken"]
+
+        # Admin delete user
+        cognito_client.admin_delete_user(UserPoolId=user_pool_id, Username=username)
+
+        # Verify user is deleted
+        with pytest.raises(ClientError) as exc_info:
+            cognito_client.admin_get_user(UserPoolId=user_pool_id, Username=username)
+        assert exc_info.value.response["Error"]["Code"] == "UserNotFoundException"
+
+    finally:
+        # Clean up resources
+        cognito_client.delete_user_pool(UserPoolId=user_pool_id)
