diff --git a/.github/workflows/MLPipelinePullRequest.yml b/.github/workflows/MLPipelinePullRequest.yml
index 30641cf..c0576c3 100644
--- a/.github/workflows/MLPipelinePullRequest.yml
+++ b/.github/workflows/MLPipelinePullRequest.yml
@@ -58,7 +58,7 @@ jobs:
       id: set_timestamp
       run: |
         echo "timestamp=$(date +%Y%m%d)" >> $GITHUB_ENV
-        echo "Generated timestamp: $timestamp"
+        echo "Generated timestamp: ${timestamp}"
 
     - name: Upload sample row dataset to S3
       env:
@@ -66,28 +66,43 @@ jobs:
         AWS_SECRET_ACCESS_KEY: ${{ secrets.ROBOT_AWS_SECRET_ACCESS_KEY }}
       run: |
         cd modules/ml-pipeline/src/pipeline/data/prepared_data/
-        aws s3 cp sample_test.parquet s3://retrofit-data-dev/sap_change_model/sample_data_for_cicd/${timestamp}_sample_test.parquet
+        aws s3 cp sample_test.parquet s3://retrofit-data-dev/sap_change_model/sample_data_for_cicd/${timestamp}/sample_test.parquet
 
     - name: Build Lambda docker Image
       run: |
         docker build . --file ./deployment/Dockerfile.prediction.lambda --tag lambda_test
 
+    - name: Run lambda docker container
+      env:
+        AWS_ACCESS_KEY_ID: ${{ secrets.ROBOT_AWS_ACCESS_KEY_ID }}
+        AWS_SECRET_ACCESS_KEY: ${{ secrets.ROBOT_AWS_SECRET_ACCESS_KEY }}
+      run: |
+        docker run -d -p 9000:8080 \
+          -e AWS_ACCESS_KEY_ID=${AWS_ACCESS_KEY_ID} \
+          -e AWS_SECRET_ACCESS_KEY=${AWS_SECRET_ACCESS_KEY} \
+          -e RUNTIME_ENVIRONMENT=dev \
+          -e PREDICTIONS_BUCKET=retrofit-sap-predictions-dev lambda_test
+
+    - name: Test Lambda endpoint
+      run: |
+        sleep 5
+        RESPONSE=$(curl -XPOST "http://localhost:9000/2015-03-31/functions/function/invocations" -d '{"body": "{\"file_location\": \"cicd://retrofit-data-dev/sap_change_model/sample_data_for_cicd/${timestamp}/sample_test.parquet\", \"property_id\": 1, \"portfolio_id\": 4, \"created_at\": \"now\"}"}')
+        echo "Lambda response: $RESPONSE"
+        if [[ $RESPONSE != *"predictions"* ]]; then
+          echo "Lambda invocation failed"
+          exit 1
+        fi
+
     - name: Remove uploaded sample row dataset from S3
       env:
         AWS_ACCESS_KEY_ID: ${{ secrets.ROBOT_AWS_ACCESS_KEY_ID }}
         AWS_SECRET_ACCESS_KEY: ${{ secrets.ROBOT_AWS_SECRET_ACCESS_KEY }}
       run: |
-        aws s3 rm s3://retrofit-data-dev/sap_change_model/sample_data_for_cicd/${timestamp}_sample_test.parquet
-    # - name: Run lambda docker container
-    #   env:
-    #     AWS_ACCESS_KEY_ID: ${{ secrets.ROBOT_AWS_ACCESS_KEY_ID }}
-    #     AWS_SECRET_ACCESS_KEY: ${{ secrets.ROBOT_AWS_SECRET_ACCESS_KEY }}
-    #   run: |
-    #     docker run -p 9000:8080 \
-    #       -e AWS_ACCESS_KEY_ID=${AWS_ACCESS_KEY_ID} \
-    #       -e AWS_SECRET_ACCESS_KEY=${AWS_SECRET_ACCESS_KEY} \
-    #       -e RUNTIME_ENVIRONMENT=dev \
-    #       -e PREDICTIONS_BUCKET=retrofit-sap-predictions-dev lambda_test
+        aws s3 rm --recursive s3://retrofit-data-dev/sap_change_model/sample_data_for_cicd/${timestamp}/
+
+    - name: Stop Lambda container
+      run: |
+        docker stop lambda_test || echo "Container already stopped"
 
   Verify-Model:
 
diff --git a/deployment/handlers/prediction_app.py b/deployment/handlers/prediction_app.py
index ac397b9..a0507bf 100644
--- a/deployment/handlers/prediction_app.py
+++ b/deployment/handlers/prediction_app.py
@@ -67,7 +67,14 @@ def handler(event, context):
 
         # TODO: Implement the loading of the model and prediction
 
-        storage_filepath = f"s3://{PREDICTIONS_BUCKET}/{portfolio_id}/{property_id}/{created_at}.parquet"
+        if body["file_location"].startswith("cicd://"):
+            storage_filepath = (
+                body["file_location"]
+                .replace("cicd://", "s3://")
+                .replace(".parquet", "_output.parquet")
+            )
+        else:
+            storage_filepath = f"s3://{PREDICTIONS_BUCKET}/{portfolio_id}/{property_id}/{created_at}.parquet"
 
         logger.info(f"--- Initiate MLModel ---")