[Python] Deflake anomaly transform tests

sdks/python/apache_beam/ml/anomaly/transforms_test.py

@@ -32,6 +32,7 @@
 
 import apache_beam as beam
 from apache_beam.ml.anomaly.aggregations import AnyVote
+from apache_beam.ml.anomaly.base import AnomalyDetector
 from apache_beam.ml.anomaly.base import AnomalyPrediction
 from apache_beam.ml.anomaly.base import AnomalyResult
 from apache_beam.ml.anomaly.base import EnsembleAnomalyDetector
@@ -116,18 +117,33 @@ def _keyed_result_is_equal_to(
   return a[0] == b[0] and _unkeyed_result_is_equal_to(a[1], b[1])
 
 
+@specifiable
+class _ValueDetector(AnomalyDetector):
+  def learn_one(self, unused_x: beam.Row) -> None:
+    pass
+
+  def score_one(self, x: beam.Row) -> float:
+    value = float(next(iter(x)))
+    return float('NaN') if value == -1 else value
+
+
 class TestAnomalyDetection(unittest.TestCase):
   class TestData:
     unkeyed_input = [
-        beam.Row(x1=1, x2=4),
-        beam.Row(x1=2, x2=4),
-        beam.Row(x1=3, x2=5),
-        beam.Row(x1=10, x2=4),  # outlier in key=1, with respect to x1
-        beam.Row(x1=2, x2=10),  # outlier in key=1, with respect to x2
-        beam.Row(x1=3, x2=4),
+        beam.Row(x1=1, x2=4, score_x1=-1.0, score_x2=-1.0),
+        beam.Row(x1=2, x2=4, score_x1=-1.0, score_x2=-1.0),
+        beam.Row(x1=3, x2=5, score_x1=2.1213203435596424, score_x2=0.0),
+        beam.Row(x1=10, x2=4, score_x1=8.0, score_x2=0.5773502691896252),
+        beam.Row(x1=2, x2=10, score_x1=0.4898979485566356, score_x2=11.5),
+        beam.Row(
+            x1=3,
+            x2=4,
+            score_x1=0.16452254913212455,
+            score_x2=0.5368754921931594),
+    ]
+    keyed_input = list(zip(itertools.repeat(1), unkeyed_input)) + [
+        (2, beam.Row(x1=100, x2=5, score_x1=-1.0, score_x2=-1.0))
     ]
-    keyed_input = list(zip(itertools.repeat(1),
-                           unkeyed_input)) + [(2, beam.Row(x1=100, x2=5))]
 
     zscore_x1_expected_predictions = [
         AnomalyPrediction(
@@ -244,10 +260,14 @@ def test_one_detector(self, input, expected):
   ])
   def test_multiple_detectors_without_aggregation(self, input, expected):
     sub_detectors = []
-    sub_detectors.append(ZScore(features=["x1"], model_id="zscore_x1"))
     sub_detectors.append(
-        ZScore(
-            features=["x2"],
+        _ValueDetector(
+            features=["score_x1"],
+            threshold_criterion=FixedThreshold(3),
+            model_id="zscore_x1"))
+    sub_detectors.append(
+        _ValueDetector(
+            features=["score_x2"],
             threshold_criterion=FixedThreshold(2),
             model_id="zscore_x2"))
 
@@ -267,10 +287,14 @@ def test_multiple_detectors_without_aggregation(self, input, expected):
   ])
   def test_multiple_sub_detectors_with_aggregation(self, input, expected):
     sub_detectors = []
-    sub_detectors.append(ZScore(features=["x1"], model_id="zscore_x1"))
     sub_detectors.append(
-        ZScore(
-            features=["x2"],
+        _ValueDetector(
+            features=["score_x1"],
+            threshold_criterion=FixedThreshold(3),
+            model_id="zscore_x1"))
+    sub_detectors.append(
+        _ValueDetector(
+            features=["score_x2"],
             threshold_criterion=FixedThreshold(2),
             model_id="zscore_x2"))
 
@@ -286,6 +310,35 @@ def test_multiple_sub_detectors_with_aggregation(self, input, expected):
       else:
         assert_that(result, equal_to(expected, _unkeyed_result_is_equal_to))
 
+  @parameterized.expand([(False, ), (True, )])
+  def test_ensemble_preserves_model_state(self, keyed):
+    # Constant values make the scores independent of processing order. Each
+    # model must produce exactly two warmup predictions per key.
+    input = [beam.Row(x1=1, x2=4)] * 6
+    counts_by_key = [(0, 6)]
+    if keyed:
+      input = [(1, row) for row in input] + [(2, beam.Row(x1=100, x2=50))] * 3
+      counts_by_key = [(1, 6), (2, 3)]
+
+    detectors = [
+        ZScore(features=[feature], model_id=feature)
+        for feature in ('x1', 'x2')
+    ]
+    expected = [(key, model_id, label) for key, count in counts_by_key
+                for model_id in ('x1', 'x2')
+                for label in [-2, -2] + [0] * (count - 2)]
+
+    with TestPipeline() as p:
+      result = (
+          p | beam.Create(input)
+          | AnomalyDetection(EnsembleAnomalyDetector(detectors)))
+      if not keyed:
+        result = result | beam.WithKeys(0)
+      labels = result | beam.FlatMap(
+          lambda item: [(item[0], prediction.model_id, prediction.label)
+                        for prediction in item[1].predictions])
+      assert_that(labels, equal_to(expected))
+
 
 class FakeNumpyModel():
   def __init__(self):