test: Added session integration tests (#527)

Author: charles-marion

docs/guide/deploy.md

@@ -164,6 +164,14 @@ GenAIChatBotStack.ApiKeysSecretNameXXXX = ApiKeysSecretName-xxxxxx
 
 **Step 11.** Login with the user created in **Step 8** and follow the instructions.
 
+**Step 12.** (Optional) Run the integration tests
+The tests require to be authenticated against your AWS Account because it will create cognito users. In addition, the tests will use `anthropic.claude-instant-v1` which needs to be enabled in Bedrock.
+
+To run the tests (Replace the url with the one you used in the steps above)
+```bash
+REACT_APP_URL=https://dxxxxxxxxxxxxx.cloudfront.net pytest integtests/
+```
+
 ## Run user interface locally
 
 To experiment with changes to the the user interface, you can run the interface locally. See the instructions in the README file of the [`lib/user-interface/react-app`](https://github.com/aws-samples/aws-genai-llm-chatbot/blob/main/lib/user-interface/react-app/README.md) folder.

integtests/chatbot-api/session_test.py

@@ -0,0 +1,78 @@
+import json 
+import uuid
+import time
+
+import pytest
+
+def test_create_session(client, default_model, default_provider, session_id):
+    request = {
+          "action": "run",
+          "modelInterface": "langchain",
+          "data": {
+            "mode": "chain",
+            "text": "test",
+            "files": [],
+            "modelName": default_model,
+            "provider": default_provider,
+            "sessionId": session_id,
+          },
+        }
+  
+    client.send_query(json.dumps(request))
+    # Need a second sessions to verify the delete all
+    request["data"]["sessionId"] = str(uuid.uuid4())
+    client.send_query(json.dumps(request))
+    
+    found = False
+    sessionFound = None
+    retries = 0
+    while not found and retries < 10:
+        time.sleep(1)
+        retries += 1
+        sessions = client.list_sessions()
+        for session in sessions:
+            if len(sessions) >= 2 and session.get("id") == session_id:
+                found = True
+                sessionFound = session
+                break
+
+    assert found == True
+    assert sessionFound.get("title") == request.get("data").get("text")
+
+def test_get_session(client, session_id, default_model):  
+    session = client.get_session(session_id)
+    assert session.get("id") == session_id
+    assert session.get("title") == "test"
+    assert len(session.get("history")) == 2
+    assert session.get("history")[0].get("type") == "human"
+    assert session.get("history")[1].get("type") == "ai"
+    
+def test_delete_session(client, session_id):  
+    session = client.delete_session(session_id)
+    assert session.get("id") == session_id
+    assert session.get("deleted") == True
+    
+    session = client.get_session(session_id)
+    assert session == None
+    
+def test_delete_user_sessions(client):  
+    sessions = client.delete_user_sessions()
+    assert len(sessions) > 0
+    assert sessions[0].get("deleted") == True
+    
+    sessions = client.list_sessions()
+    retries = 0
+    while True:
+        time.sleep(1)
+        retries += 1
+        if retries > 10:
+            pytest.fail()
+        elif len(client.list_sessions()) == 0:
+            break
+    
+@pytest.fixture(scope="package")
+def session_id():
+    return str(uuid.uuid4())
+    
+
+

integtests/clients/appsync_client.py

@@ -0,0 +1,88 @@
+import os
+from gql import Client
+from gql.transport.aiohttp import AIOHTTPTransport
+from gql.dsl import DSLMutation, DSLSchema, DSLQuery, dsl_gql
+from graphql import print_ast
+
+
+class AppSyncClient:
+    def __init__(self, endpoint: str, id_token: str) -> None:
+        self.transport = AIOHTTPTransport(
+            url=endpoint
+        )
+        if id_token is not None:
+            self.transport.headers = {"Authorization": id_token}
+        dir_path = os.path.dirname(os.path.realpath(__file__))
+        with open(dir_path + "/../../lib/chatbot-api/schema/schema.graphql") as f:
+            schema_string = f.read()
+
+        # remove AWS specific syntax from the schema
+        schema_string = schema_string.replace("@aws_cognito_user_pools", "")
+        schema_string = schema_string.replace("@aws_iam", "")
+        schema_string = schema_string.replace(
+            '@aws_subscribe(mutations: ["publishResponse"])', ""
+        )
+        schema_string = schema_string.replace("AWSDateTime", "String")
+
+        self.client = Client(transport=self.transport, schema=schema_string)
+        self.schema = DSLSchema(self.client.schema)
+
+    def send_query(self, data: str):
+        query = dsl_gql(DSLMutation(self.schema.Mutation.sendQuery.args(data=data)))
+        return self.client.execute(query)
+
+    def delete_session(self, id: str):
+        query = dsl_gql(
+            DSLMutation(
+                self.schema.Mutation.deleteSession.args(id=id).select(
+                    self.schema.DeleteSessionResult.id,
+                    self.schema.DeleteSessionResult.deleted,
+                )
+            )
+        )
+        return self.client.execute(query).get("deleteSession")
+
+    def delete_user_sessions(self):
+        query = dsl_gql(
+            DSLMutation(
+                self.schema.Mutation.deleteUserSessions.select(
+                    self.schema.DeleteSessionResult.id,
+                    self.schema.DeleteSessionResult.deleted,
+                )
+            )
+        )
+        return self.client.execute(query).get("deleteUserSessions")
+
+    def list_sessions(self):
+        query = dsl_gql(
+            DSLQuery(
+                self.schema.Query.listSessions.select(
+                    self.schema.Session.id,
+                    self.schema.Session.startTime,
+                    self.schema.Session.title,
+                    self.schema.Session.history.select(
+                        self.schema.SessionHistoryItem.type,
+                        self.schema.SessionHistoryItem.content,
+                        self.schema.SessionHistoryItem.metadata,
+                    ),
+                )
+            )
+        )
+        return self.client.execute(query).get("listSessions")
+
+    def get_session(self, id):
+        query = dsl_gql(
+            DSLQuery(
+                self.schema.Query.getSession(id=id).select(
+                    self.schema.Session.id,
+                    self.schema.Session.startTime,
+                    self.schema.Session.title,
+                    self.schema.Session.history.select(
+                        self.schema.SessionHistoryItem.type,
+                        self.schema.SessionHistoryItem.content,
+                        self.schema.SessionHistoryItem.metadata,
+                    ),
+                )
+            )
+        )
+        return self.client.execute(query).get("getSession")

integtests/clients/cognito_client.py

@@ -0,0 +1,61 @@
+
+import boto3
+import string
+import random
+
+class CognitoClient:
+    def __init__(self, region: str, user_pool_id: str, client_id: str) -> None:
+        self.user_pool_id = user_pool_id
+        self.client_id = client_id
+        self.cognito_idp_client = boto3.client('cognito-idp', region_name=region)
+
+    def get_token(self, email: str) -> None:
+        try:
+            self.cognito_idp_client.admin_get_user(
+                UserPoolId=self.user_pool_id,
+                Username=email,
+            )
+        except self.cognito_idp_client.exceptions.UserNotFoundException:
+            self.cognito_idp_client.admin_create_user(
+                UserPoolId=self.user_pool_id,
+                Username=email,
+                UserAttributes=[
+                    {
+                        'Name': 'email',
+                        'Value': email
+                    },
+                    {
+                        'Name': 'email_verified',
+                        'Value': 'True'
+                    }
+                ],
+                MessageAction="SUPPRESS",
+            )
+            
+        password = self.get_password()
+        self.cognito_idp_client.admin_set_user_password(
+            UserPoolId=self.user_pool_id,
+            Username=email,
+            Password=password,
+            Permanent=True
+        )
+        
+        response = self.cognito_idp_client.admin_initiate_auth(
+            UserPoolId=self.user_pool_id,
+            ClientId=self.client_id,
+            AuthFlow="ADMIN_NO_SRP_AUTH",
+            AuthParameters={
+                "USERNAME": email,
+                "PASSWORD": password
+            }
+        )
+        
+        return response["AuthenticationResult"]["IdToken"]
+        
+    def get_password(self):
+        return "".join(
+            random.choices(string.ascii_uppercase, k=10) +
+            random.choices(string.ascii_lowercase, k=10) +
+            random.choices(string.digits, k=5) +
+            random.choices(string.punctuation, k=3)
+        )

integtests/conftest.py

@@ -0,0 +1,41 @@
+import json
+import os
+
+from urllib.request import urlopen
+
+import pytest
+
+from clients.cognito_client import CognitoClient
+from clients.appsync_client import AppSyncClient
+
+@pytest.fixture(scope="session")
+def client(config):
+    user_pool_id = config.get("aws_user_pools_id")
+    region = config.get("aws_cognito_region")
+    user_pool_client_id = config.get("aws_user_pools_web_client_id")
+    endpoint = config.get("aws_appsync_graphqlEndpoint")
+    
+    cognito = CognitoClient(region=region, user_pool_id=user_pool_id, client_id=user_pool_client_id)
+    email = "integ-test-user@example.local"
+    
+    return AppSyncClient(endpoint=endpoint, id_token=cognito.get_token(email=email))
+
+@pytest.fixture(scope="session")
+def unauthenticated_client(config):    
+    endpoint = config.get("aws_appsync_graphqlEndpoint")
+    return AppSyncClient(endpoint=endpoint, id_token=None)
+
+@pytest.fixture(scope="session")
+def default_model():
+    return "anthropic.claude-instant-v1"
+
+@pytest.fixture(scope="session")
+def default_provider():
+    return "bedrock"
+
+@pytest.fixture(scope="session")
+def config():
+    if "REACT_APP_URL" not in os.environ:
+        raise IndexError("Please set the environment variable REACT_APP_URL")
+    response = urlopen(os.environ['REACT_APP_URL'] + "/aws-exports.json")
+    return json.loads(response.read())

integtests/security/unauthorized_tet.py

@@ -0,0 +1,16 @@
+
+import pytest
+from gql.transport.exceptions import TransportQueryError
+
+def test_unauthenticated(unauthenticated_client): 
+    match = "UnauthorizedException" 
+    with pytest.raises(TransportQueryError, match=match):
+        unauthenticated_client.send_query("")
+    with pytest.raises(TransportQueryError, match=match):
+        unauthenticated_client.get_session("id")
+    with pytest.raises(TransportQueryError, match=match):
+        unauthenticated_client.list_sessions()
+    with pytest.raises(TransportQueryError, match=match):
+        unauthenticated_client.delete_session("id")
+    with pytest.raises(TransportQueryError, match=match):
+        unauthenticated_client.delete_user_sessions()

pytest_requirements.txt

@@ -1,4 +1,5 @@
 pytest==7.4.0
 pytest-mock==3.11.1
+gql==3.4.1
 aws_lambda_powertools==2.42.0
 -r lib/shared/layers/common/requirements.txt