fix(ta-gql): get testsuites

graphql_api/tests/test_test_analytics.py

@@ -744,7 +744,41 @@ def test_gql_query_flake_aggregates(self, repository, store_in_redis):
             "flakeCount": 1,
         }
 
-    def test_gql_query_with_new_ta(self, mocker, repository, snapshot):
+    def test_gql_query_test_suites(self, repository, store_in_redis):
+        query = base_gql_query % (
+            repository.author.username,
+            repository.name,
+            """
+            testSuites
+            """,
+        )
+
+        result = self.gql_request(query, owner=repository.author)
+
+        assert sorted(result["owner"]["repository"]["testAnalytics"]["testSuites"]) == [
+            "testsuite1",
+            "testsuite2",
+            "testsuite3",
+            "testsuite4",
+            "testsuite5",
+        ]
+
+    def test_gql_query_test_suites_term(self, repository, store_in_redis):
+        query = base_gql_query % (
+            repository.author.username,
+            repository.name,
+            """
+            testSuites(term: "testsuite1")
+            """,
+        )
+
+        result = self.gql_request(query, owner=repository.author)
+
+        assert result["owner"]["repository"]["testAnalytics"]["testSuites"] == [
+            "testsuite1",
+        ]
+
+    def test_gql_query_with_new_ta_v1(self, mocker, repository, snapshot):
         # set the feature flag
         mocker.patch("rollouts.READ_NEW_TA.check_value", return_value=True)
 
@@ -757,7 +791,8 @@ def test_gql_query_with_new_ta(self, mocker, repository, snapshot):
         storage.write_file(
             settings.GCS_BUCKET_NAME,
             f"test_analytics/branch_rollups/{repository.repoid}/{repository.branch}.arrow",
-            test_results_table_no_version.write_ipc(None).getvalue(),
+            test_results_table_v1.write_ipc(None).getvalue(),
+            metadata={"version": "1"},
         )
 
         # run the GQL query
@@ -784,6 +819,18 @@ def test_gql_query_with_new_ta(self, mocker, repository, snapshot):
                     }
                 }
             }
+            flakeAggregates {
+                flakeRate
+                flakeCount
+            }
+            testResultsAggregates {
+                totalDuration
+                slowestTestsDuration
+                totalFails
+                totalSkips
+                totalSlowTests
+            }
+            testSuites
             """,
         )
 
@@ -804,6 +851,28 @@ def test_gql_query_with_new_ta(self, mocker, repository, snapshot):
             ]
         ]
 
+        assert sorted(result["owner"]["repository"]["testAnalytics"]["testSuites"]) == [
+            "testsuite0",
+            "testsuite1",
+            "testsuite2",
+            "testsuite3",
+            "testsuite4",
+        ]
+
+        assert result["owner"]["repository"]["testAnalytics"]["flakeAggregates"] == {
+            "flakeRate": (1 / 15),
+            "flakeCount": 1,
+        }
+
+        assert result["owner"]["repository"]["testAnalytics"][
+            "testResultsAggregates"
+        ] == {
+            "totalDuration": 7500.0,
+            "slowestTestsDuration": 2500.0,
+            "totalFails": 50,
+            "totalSkips": 25,
+            "totalSlowTests": 1,
+        }
         storage.delete_file(
             settings.GCS_BUCKET_NAME,
             f"test_analytics/branch_rollups/{repository.repoid}/{repository.branch}.arrow",

graphql_api/types/test_analytics/test_analytics.py

@@ -312,7 +312,7 @@ def get_test_suites(
     if table is None:
         return []
 
-    testsuites = table.select(pl.col("testsuite")).unique()
+    testsuites = table.select(pl.col("testsuite").explode()).unique()
 
     if term:
         testsuites = testsuites.filter(pl.col("testsuite").str.starts_with(term))

graphql_api/types/test_results_aggregates/test_results_aggregates.py

@@ -24,7 +24,7 @@ class TestResultsAggregates:
 
 
 def calculate_aggregates(table: pl.DataFrame) -> pl.DataFrame:
-    return table.select(
+    return table.with_columns(pl.col("avg_duration").fill_nan(0)).select(
         (
             pl.col("avg_duration")
             * (pl.col("total_pass_count") + pl.col("total_fail_count"))

utils/test_results.py

@@ -198,7 +198,7 @@ def v1_agg_table(table: pl.LazyFrame) -> pl.LazyFrame:
         pl.col("avg_duration") * (pl.col("pass_count") + pl.col("fail_count"))
     ).sum() / (pl.col("pass_count") + pl.col("fail_count")).sum()
 
-    table = table.group_by("computed_name").agg(
+    table = table.group_by(pl.col("computed_name").alias("name")).agg(
         pl.col("testsuite").alias(
             "testsuite"
         ),  # TODO: filter by this before we aggregate