From d3f77fabb83b8904aa6654818574b6d90f8d3b1a Mon Sep 17 00:00:00 2001
From: Danny Zhao <zikangz@amazon.com>
Date: Wed, 13 Sep 2023 22:56:53 -0700
Subject: [PATCH] Fix AnomalyDetectionMathExpression

---
 .../metric/AnomalyDetectionMathExpression.ts   | 18 ++++++++++++++++--
 .../CustomMonitoring.test.ts.snap              |  6 +++---
 2 files changed, 19 insertions(+), 5 deletions(-)

diff --git a/lib/common/metric/AnomalyDetectionMathExpression.ts b/lib/common/metric/AnomalyDetectionMathExpression.ts
index a3b94feb..4e6f3eca 100644
--- a/lib/common/metric/AnomalyDetectionMathExpression.ts
+++ b/lib/common/metric/AnomalyDetectionMathExpression.ts
@@ -31,15 +31,29 @@ export class AnomalyDetectionMathExpression extends MathExpression {
   createAlarm(scope: Construct, id: string, props: CreateAlarmOptions): Alarm {
     const alarm = super.createAlarm(scope, id, props);
 
+    // `usingMetrics` of an anomaly detection alarm can only ever have one entry.
+    // Should the entry be a math expression, the math expression can have its own `usingMetrics`.
+    const finalExpressionId = Object.keys(this.usingMetrics)[0];
+
     // https://github.com/aws/aws-cdk/issues/10540#issuecomment-725222564
     const cfnAlarm = alarm.node.defaultChild as CfnAlarm;
     cfnAlarm.addPropertyDeletionOverride("Threshold");
     (cfnAlarm.metrics as CfnAlarm.MetricDataQueryProperty[]).forEach(
       (metric, index) => {
-        if (metric.expression) {
+        // To create an anomaly detection alarm, returned data should be set to true on two MetricDataQueryProperty(s):
+        // 1. The metric or math expression that is being evaluated for anomaly detection (eg. expr_1)
+        // 2. The actual expression of anomaly detection (eg. ANOMALY_DETECTION_BAND(expr_1, 1))
+        let returnData = false;
+
+        if (metric.expression?.includes("ANOMALY_DETECTION_BAND")) {
+          // thresholdMetricId is the ID of the ANOMALY_DETECTION_BAND function used as the threshold for the alarm.
           cfnAlarm.thresholdMetricId = metric.id;
+          returnData = true;
+        } else if (metric.id === finalExpressionId) {
+          returnData = true;
         }
-        cfnAlarm.addPropertyOverride(`Metrics.${index}.ReturnData`, true);
+
+        cfnAlarm.addPropertyOverride(`Metrics.${index}.ReturnData`, returnData);
       }
     );
 
diff --git a/test/monitoring/custom/__snapshots__/CustomMonitoring.test.ts.snap b/test/monitoring/custom/__snapshots__/CustomMonitoring.test.ts.snap
index 307b7f7a..975ebb96 100644
--- a/test/monitoring/custom/__snapshots__/CustomMonitoring.test.ts.snap
+++ b/test/monitoring/custom/__snapshots__/CustomMonitoring.test.ts.snap
@@ -422,7 +422,7 @@ Object {
               "Period": 300,
               "Stat": "Average",
             },
-            "ReturnData": true,
+            "ReturnData": false,
           },
           Object {
             "Id": "m2",
@@ -434,10 +434,10 @@ Object {
               "Period": 300,
               "Stat": "Average",
             },
-            "ReturnData": true,
+            "ReturnData": false,
           },
         ],
-        "ThresholdMetricId": "alarm_39cd98a4dd0f0",
+        "ThresholdMetricId": "expr_1",
         "TreatMissingData": "missing",
       },
       "Type": "AWS::CloudWatch::Alarm",