Add step_test_data to luci_analysis.query_stability

Bug: 346598710
Change-Id: I6f1dccf22b651b45cb8b0c7330077b0e3e6f4e40

Recipe-Manual-Change: chromiumos
Change-Id: I6f1dccf22b651b45cb8b0c7330077b0e3e6f4e40
Reviewed-on: https://chromium-review.googlesource.com/c/infra/luci/recipes-py/+/5738337
Reviewed-by: Patrick Meiring <meiring@google.com>
Commit-Queue: Stephanie Kim <kimstephanie@google.com>
diff --git a/README.recipes.md b/README.recipes.md
index b396ca3..fcb56dc 100644
--- a/README.recipes.md
+++ b/README.recipes.md
@@ -2648,7 +2648,7 @@
 
 #### **class [LuciAnalysisApi](/recipe_modules/luci_analysis/api.py#30)([RecipeApi](/recipe_engine/recipe_api.py#433)):**
 
-&mdash; **def [lookup\_bug](/recipe_modules/luci_analysis/api.py#263)(self, bug_id, system='monorail'):**
+&mdash; **def [lookup\_bug](/recipe_modules/luci_analysis/api.py#273)(self, bug_id, system='monorail'):**
 
 Looks up the rule associated with a given bug.
 
@@ -2664,7 +2664,7 @@
 Returns:
   list of rules (str), Format: projects/{project}/rules/{rule_id}
 
-&mdash; **def [query\_cluster\_failures](/recipe_modules/luci_analysis/api.py#298)(self, cluster_name):**
+&mdash; **def [query\_cluster\_failures](/recipe_modules/luci_analysis/api.py#308)(self, cluster_name):**
 
 Queries examples of failures in the given cluster.
 
@@ -2691,7 +2691,7 @@
 Returns:
   List of TestVariantFailureRateAnalysis protos
 
-&mdash; **def [query\_stability](/recipe_modules/luci_analysis/api.py#117)(self, test_variant_position_list, project='chromium'):**
+&mdash; **def [query\_stability](/recipe_modules/luci_analysis/api.py#125)(self, test_variant_position_list, project='chromium'):**
 
 Queries LUCI Analysis for test stability.
 
@@ -2704,7 +2704,7 @@
 Raises:
   StepFailure if query is invalid or service returns unexpected responses.
 
-&mdash; **def [query\_test\_history](/recipe_modules/luci_analysis/api.py#153)(self, test_id, project='chromium', sub_realm=None, variant_predicate=None, partition_time_range=None, submitted_filter=None, page_size=1000, page_token=None):**
+&mdash; **def [query\_test\_history](/recipe_modules/luci_analysis/api.py#163)(self, test_id, project='chromium', sub_realm=None, variant_predicate=None, partition_time_range=None, submitted_filter=None, page_size=1000, page_token=None):**
 
 A wrapper method to use `luci.analysis.v1.TestHistory` `Query` API.
 
@@ -2735,7 +2735,7 @@
 Returns:
   (list of parsed luci.analysis.v1.TestVerdict objects, next page token)
 
-&mdash; **def [query\_variants](/recipe_modules/luci_analysis/api.py#214)(self, test_id, project='chromium', sub_realm=None, variant_predicate=None, page_size=1000, page_token=None):**
+&mdash; **def [query\_variants](/recipe_modules/luci_analysis/api.py#224)(self, test_id, project='chromium', sub_realm=None, variant_predicate=None, page_size=1000, page_token=None):**
 
 A wrapper method to use `luci.analysis.v1.TestHistory` `QueryVariants`
 API.
@@ -2762,7 +2762,7 @@
   (list of VariantInfo { variant_hash: str, variant: { def: dict } },
    next page token)
 
-&mdash; **def [rule\_name\_to\_cluster\_name](/recipe_modules/luci_analysis/api.py#287)(self, rule):**
+&mdash; **def [rule\_name\_to\_cluster\_name](/recipe_modules/luci_analysis/api.py#297)(self, rule):**
 
 Convert the resource name for a rule to its corresponding cluster.
 Args:
diff --git a/recipe_modules/luci_analysis/api.py b/recipe_modules/luci_analysis/api.py
index c83cf4e..2cd87e8 100644
--- a/recipe_modules/luci_analysis/api.py
+++ b/recipe_modules/luci_analysis/api.py
@@ -114,6 +114,14 @@
           for d in failure_analysis_dicts
       ]
 
+  def _query_stability_step_test_data(self, test_ids):
+    stability_analyses = [
+        self.test_api.generate_stability_analysis(test_id)
+        for test_id in test_ids
+    ]
+    return self.m.json.test_api.output_stream(
+        self.test_api.generate_stability_response(stability_analyses))
+
   def query_stability(self, test_variant_position_list, project='chromium'):
     """Queries LUCI Analysis for test stability.
 
@@ -134,6 +142,8 @@
               'project': project,
               'testVariants': test_variant_position_list,
           },
+          step_test_data=lambda: self._query_stability_step_test_data(
+              [t['testId'] for t in test_variant_position_list]),
       )
       # This is not likely to happen, although invalid request may result in a
       # StepFailure directly raised from the above step.