Update docstrings, docs, and typing

Author: mbeacom

.DS_Store

@@ -90,16 +90,12 @@ def login(self, username="", password=""):
             _xsrf_token (str): The XSRF token to be used for subsequent API requests.
 
         """
-        endpoint: str = "login"
         _username: str = self.config.active_config["username"] or username
         _password: str = self.config.active_config["password"] or password
         _auth: Dict[str, str] = {"username": _username, "password": _password}
 
         # Attempt to login to the CloudEndure API via a POST request.
-        response: requests.Response = self.api_call(
-            "login", "post", data=json.dumps(_auth)
-        )
-        # response: requests.Response = self.session.post(f'{self.api_endpoint}/{endpoint}', json=_auth)
+        response: requests.Response = self.api_call("login", "post", data=json.dumps(_auth))
 
         # Check whether or not the request was successful.
         if response.status_code not in [200, 307]:
@@ -117,8 +113,11 @@ def login(self, username="", password=""):
                 )
             raise CloudEndureUnauthorized()
 
-        # print('response: ', response, response.cookies)
-        _xsrf_token: str = str(response.cookies["XSRF-TOKEN"])
+        # Grab the XSRF token received from the response, as stored in cookies.
+        # _xsrf_token: str = str(response.cookies["XSRF-TOKEN"])
+        _xsrf_token: str = str(response.cookies.get("XSRF-TOKEN", ""))
+        if not _xsrf_token:
+            raise CloudEndureException("Failed to fetch a token from CloudEndure!")
 
         # Strip the XSRF token of wrapping double-quotes from the cookie.
         if _xsrf_token.startswith('"') and _xsrf_token.endswith('"'):

cloudendure/api.py

@@ -30,6 +30,7 @@
 MIGRATION_WAVE: str = os.environ.get("CLOUDENDURE_MIGRATION_WAVE", "0")
 CLONE_STATUS: str = os.environ.get("CLOUDENDURE_CLONE_STATUS", "NOT_STARTED")
 MAX_LAG_TTL: int = int(os.environ.get("CLOUDENDURE_MAX_LAG_TTL", "90"))
+SHARE_IMAGE: str = os.environ.get("", "")
 
 
 class CloudEndure:
@@ -201,6 +202,7 @@ def update_blueprint(
         except Exception as e:
             print(f"Updating blueprint task failed! {e}")
             return False
+        return True
 
     def launch(self, project_id=global_project_id, launch_type="test", dry_run=False):
         """Launch the test target instances."""

cloudendure/cloudendure.py

@@ -0,0 +1,20 @@
+# -*- coding: utf-8 -*-
+"""Define the Lambda specific exceptions."""
+
+
+class LambdaException(Exception):
+    """Define the generic AWS Lambda exception."""
+
+    pass
+
+
+class InvalidPayload(LambdaException):
+    """Define the exception to be raised when an invalid payload is encountered."""
+
+    pass
+
+
+class ImproperlyConfigured(LambdaException):
+    """Define the exception to be raised if the environment is improperly configured or missing."""
+
+    pass

lambda/exceptions.py

@@ -4,17 +4,20 @@
 Attributes:
     LOG_LEVEL (str):
     REGION_OVERRIDE (str): If provided, this value will override the default AWS region.
-    logger (logging.Logger):
+    logger (logging.Logger): The logger to be used throughout execution of the AWS Lambda.
 
 """
 import datetime
 import json
 import logging
 import os
+from typing import Any, Dict, List
 
 import boto3
 from botocore.exceptions import ClientError
 
+from .exceptions import ImproperlyConfigured, InvalidPayload
+
 LOG_LEVEL = getattr(logging, os.environ.get("CLOUDENDURE_LOGLEVEL", "INFO"))
 REGION_OVERRIDE = os.environ.get("CLOUDENDURE_REGION_OVERRIDE", "")
 
@@ -23,23 +26,51 @@
 logger.setLevel(LOG_LEVEL)
 
 
-def send_sqs_message(imageinfo):
-    """Sends a SQS message with the AMI information
-    that was created from the migrated instance
-    that passed testing post migration in CloudEndure
+def send_sqs_message(image_info: Dict) -> bool:
+    """Send a SQS message.
+
+    The message includes the AMI information that was created from the migrated
+    instance that passed testing post migration in CloudEndure.
+
+    Raises:
+        ClientError: The exception is raised in the event of a boto3 client error.
+        ImproperlyConfigured: The exception is raised in the event of missing or invalid
+            environment configuration settings.
+
+    Returns:
+        bool:  Whether or not the message has been sent successfully.
+
     """
+    queue_url = os.environ.get("QueueURL", "")
+    if not queue_url:
+        raise ImproperlyConfigured("Missing environment variable:  QueueURL")
+
     try:
-        message = json.dumps(imageinfo)
+        message = json.dumps(image_info)
         sqsclient = boto3.client("sqs")
-        sqsclient.send_message(QueueUrl=os.environ["QueueURL"], MessageBody=message)
-
-    except ClientError as err:
-        logger.error(err.response)
+        sqsclient.send_message(QueueUrl=queue_url, MessageBody=message)
+    except ClientError as e:
+        logger.error(e.response)
+        return False
+    except ImproperlyConfigured as e:
+        logger.error(e)
+        return False
+    return True
 
 
-def create_ami(project_id, instance_id):
+def create_ami(project_id: str, instance_id: str) -> bool:
     """Create an AMI from the specified instance.
 
+    Args:
+        project_id (str): The ID associated with the Project.
+        instance_id (str): The ID associated with the AWS instance.
+
+    Raises:
+        ClientError: The exception is raised in the event of a boto3 client error.
+
+    Returns:
+        bool: Whether or not the AMI has been created successfully.
+
     """
     try:
         _ec2_client = boto3.client("ec2")
@@ -53,7 +84,10 @@ def create_ami(project_id, instance_id):
             NoReboot=True,
         )
         logger.info("AMI Id: %s", ec2_image)
-        _filters = [{"Name": "resource-id", "Values": [instance_id]}]
+        _filters: List[Dict] = [{
+            "Name": "resource-id",
+            "Values": [instance_id],
+        }]
 
         # Tag the newly created AMI by getting the tags of the migrated instance to copy to the AMI.
         ec2_tags = _ec2_client.describe_tags(Filters=_filters)
@@ -69,26 +103,46 @@ def create_ami(project_id, instance_id):
 
     except ClientError as err:
         logger.error(err.response)
+        return False
+    return True
 
 
-def lambda_handler(event, context):
-    """Lambda entry point"""
-    logger.info(event)
+def lambda_handler(event: Dict[str, Any], context: Dict[str, Any]) -> bool:
+    """Define the AWS Lambda entry point and handler.
 
-    try:
-        json_sns_message = json.loads(event["Records"][0]["Sns"]["Message"])
+    Args:
+        event (str): The event performed against Lambda.
+        context (dict): The context of the request performed against Lambda.
 
-        if json_sns_message["Pass"] != "True":
-            logger.error(
-                "%s did not pass post migration testing! Not creating an AMI."
-                % (json_sns_message["instanceId"])
-            )
+    Raises:
+        ClientError: The exception is raised in the event of a boto3 client error.
+        InvalidPayload: The exception is raised in the event of an invalid payload.
+
+    Returns:
+        bool: Whether or not the lambda function has executed successfully.
+
+    """
+    logger.debug(event)
+
+    event_records = event.get("Records", [])
+    if not event_records:
+        return False
 
+    try:
+        event_message = event_records[0].get("Sns", {}).get("Message", "")
+        json_sns_message = json.loads(event_message)
+        instance_id: str = json_sns_message.get("instanceId", "")
+        project_id: str = json_sns_message.get("projectId", "")
+
+        if json_sns_message.get("Pass", "NA") != "True":
+            raise InvalidPayload(f"{instance_id} did not pass post migration testing! Not creating an AMI.")
         else:
-            logger.info(
-                "%s passed post migration testing. Creating an AMI."
-                % (json_sns_message["instanceId"])
-            )
-            create_ami("", json_sns_message["instanceId"])
-    except ClientError as err:
-        logger.error(err.response)
+            logger.info("%s passed post migration testing. Creating an AMI." % (instance_id))
+            create_ami(project_id, instance_id)
+    except ClientError as e:
+        logger.error(e.response)
+        return False
+    except InvalidPayload as e:
+        logger.error(e)
+        return False
+    return True

lambda/handler.py

@@ -15,6 +15,12 @@ generate:
     - cloudendure.models++
   - code/cloudendure/utils.md:
     - cloudendure.utils++
+  - code/lambda/copy_ami.md:
+    - lambda.copy_ami++
+  - code/lambda/exceptions.md:
+    - lambda.exceptions++
+  - code/lambda/handler.md:
+    - lambda.handler++
 
 pages:
   - General:
@@ -28,6 +34,10 @@ pages:
       - Exceptions: code/cloudendure/exceptions.md
       - Models: code/cloudendure/models.md
       - Utilities: code/cloudendure/utils.md
+    - AWS Lambda:
+      - Handler: code/lambda/handler.md
+      - Exceptions: code/lambda/exceptions.md
+      - Copy AMI Utils: code/lambda/copy_ami.md
   - API Documentation:
     - MAIN: api/API_README.md << docs/API_README.md
     - AccountApi: api/AccountApi.md << docs/AccountApi.md
@@ -133,3 +143,4 @@ theme:
 
 additional_search_paths:
   - cloudendure/
+  - lambda/