[SDTEST-2993] Switch to paginated endpoint for retrieving known tests

tracer/src/Datadog.Trace/Ci/Net/TestOptimizationClient.GetKnownTestsAsync.cs

@@ -1,4 +1,4 @@
-// <copyright file="TestOptimizationClient.GetKnownTestsAsync.cs" company="Datadog">
+// <copyright file="TestOptimizationClient.GetKnownTestsAsync.cs" company="Datadog">
 // Unless explicitly stated otherwise all files in this repository are licensed under the Apache 2 License.
 // This product includes software developed at Datadog (https://www.datadoghq.com/). Copyright 2017 Datadog, Inc.
 // </copyright>
@@ -10,6 +10,7 @@
 using Datadog.Trace.Ci.Telemetry;
 using Datadog.Trace.Telemetry;
 using Datadog.Trace.Telemetry.Metrics;
+using Datadog.Trace.Util;
 using Datadog.Trace.Util.Json;
 using Datadog.Trace.Vendors.Newtonsoft.Json;
 
@@ -23,6 +24,7 @@ internal sealed partial class TestOptimizationClient
 {
     private const string KnownTestsUrlPath = "api/v2/ci/libraries/tests";
     private const string KnownTestsType = "ci_app_libraries_tests_request";
+    private const int MaxKnownTestsPages = 10_000;
     private Uri? _knownTestsUrl;
 
     public async Task<KnownTestsResponse> GetKnownTestsAsync()
@@ -34,36 +36,81 @@ public async Task<KnownTestsResponse> GetKnownTestsAsync()
         }
 
         _knownTestsUrl ??= GetUriFromPath(KnownTestsUrlPath);
-        var query = new DataEnvelope<Data<KnownTestsQuery>>(
-            new Data<KnownTestsQuery>(
-                _commitSha,
-                KnownTestsType,
-                new KnownTestsQuery(_serviceName, _environment, _repositoryUrl, GetTestConfigurations())),
-            null);
 
-        var jsonQuery = JsonHelper.SerializeObject(query, SerializerSettings);
-        Log.Debug("TestOptimizationClient: KnownTests.JSON RQ = {Json}", jsonQuery);
+        var configurations = GetTestConfigurations();
+        KnownTestsResponse.KnownTestsModules? aggregateTests = null;
+        string? pageState = null;
+        var pageNumber = 0;
 
-        string? queryResponse;
-        try
+        do
         {
-            queryResponse = await SendJsonRequestAsync<KnownTestsCallbacks>(_knownTestsUrl, jsonQuery).ConfigureAwait(false);
-        }
-        catch (Exception ex)
-        {
-            TelemetryFactory.Metrics.RecordCountCIVisibilityKnownTestsRequestErrors(MetricTags.CIVisibilityErrorType.Network);
-            Log.Error(ex, "TestOptimizationClient: Known tests request failed.");
-            throw;
+            pageNumber++;
+            var query = new DataEnvelope<Data<KnownTestsQuery>>(
+                new Data<KnownTestsQuery>(
+                    _commitSha,
+                    KnownTestsType,
+                    new KnownTestsQuery(_serviceName, _environment, _repositoryUrl, configurations, new PageInfoRequest(pageState))),
+                null);
+
+            var jsonQuery = JsonHelper.SerializeObject(query, SerializerSettings);
+            Log.Debug("TestOptimizationClient: KnownTests.JSON RQ (page {PageNumber}) = {Json}", pageNumber, jsonQuery);
+
+            string? queryResponse;
+            try
+            {
+                queryResponse = await SendJsonRequestAsync<KnownTestsCallbacks>(_knownTestsUrl, jsonQuery).ConfigureAwait(false);
+            }
+            catch (Exception ex)
+            {
+                TelemetryFactory.Metrics.RecordCountCIVisibilityKnownTestsRequestErrors(MetricTags.CIVisibilityErrorType.Network);
+                Log.Error<int>(ex, "TestOptimizationClient: Known tests request failed on page {PageNumber}.", pageNumber);
+                throw;
+            }
+
+            Log.Debug("TestOptimizationClient: KnownTests.JSON RS (page {PageNumber}) = {Json}", pageNumber, queryResponse);
+            if (StringUtil.IsNullOrEmpty(queryResponse))
+            {
+                break;
+            }
+
+            var deserializedResult = JsonHelper.DeserializeObject<DataEnvelope<Data<KnownTestsPageResponse>?>>(queryResponse);
+            var pageResponse = deserializedResult.Data?.Attributes;
+
+            if (pageResponse is null)
+            {
+                break;
+            }
+
+            var page = pageResponse.Value;
+
+            // Merge page tests into aggregate
+            MergeKnownTests(ref aggregateTests, page.Tests);
+
+            // Check pagination
+            var pageInfo = page.PageInfo;
+            if (pageInfo is not { HasNext: true })
+            {
+                // No page_info or has_next is false — we're done
+                break;
+            }
+
+            var cursor = pageInfo.Value.Cursor;
+            if (StringUtil.IsNullOrEmpty(cursor))
+            {
+                Log.Warning<int>("TestOptimizationClient: Known tests response has has_next=true but no cursor on page {PageNumber}. Aborting pagination.", pageNumber);
+                return default;
+            }
+
+            pageState = cursor;
         }
+        while (pageNumber < MaxKnownTestsPages);
 
-        Log.Debug("TestOptimizationClient: KnownTests.JSON RS = {Json}", queryResponse);
-        if (string.IsNullOrEmpty(queryResponse))
+        if (pageNumber >= MaxKnownTestsPages)
         {
-            return default;
+            Log.Warning<int>("TestOptimizationClient: Known tests pagination exceeded maximum of {MaxPages} pages. Returning data collected so far.", MaxKnownTestsPages);
         }
 
-        var deserializedResult = JsonHelper.DeserializeObject<DataEnvelope<Data<KnownTestsResponse>?>>(queryResponse);
-        var finalResponse = deserializedResult.Data?.Attributes ?? default;
+        var finalResponse = new KnownTestsResponse(aggregateTests);
 
         // Count the number of tests for telemetry
         var testsCount = 0;
@@ -75,7 +122,7 @@ public async Task<KnownTestsResponse> GetKnownTestsAsync()
                 {
                     foreach (var testsArray in suitesDictionary.Values)
                     {
-                        testsCount += testsArray?.Length ?? 0;
+                        testsCount += testsArray?.Count ?? 0;
                     }
                 }
             }
@@ -85,6 +132,52 @@ public async Task<KnownTestsResponse> GetKnownTestsAsync()
         return finalResponse;
     }
 
+    private static void MergeKnownTests(ref KnownTestsResponse.KnownTestsModules? aggregate, KnownTestsResponse.KnownTestsModules? page)
+    {
+        if (page is null)
+        {
+            return;
+        }
+
+        aggregate ??= new KnownTestsResponse.KnownTestsModules();
+
+        if (page.Count == 0)
+        {
+            return;
+        }
+
+        foreach (var moduleEntry in page)
+        {
+            if (moduleEntry.Value is null)
+            {
+                continue;
+            }
+
+            if (!aggregate.TryGetValue(moduleEntry.Key, out var existingSuites) || existingSuites is null)
+            {
+                existingSuites = new KnownTestsResponse.KnownTestsSuites();
+                aggregate[moduleEntry.Key] = existingSuites;
+            }
+
+            foreach (var suiteEntry in moduleEntry.Value)
+            {
+                if (suiteEntry.Value is null or { Count: 0 })
+                {
+                    continue;
+                }
+
+                if (!existingSuites.TryGetValue(suiteEntry.Key, out var existingTests) || existingTests is null)
+                {
+                    existingSuites[suiteEntry.Key] = suiteEntry.Value;
+                }
+                else
+                {
+                    existingTests.AddRange(suiteEntry.Value);
+                }
+            }
+        }
+    }
+
     private readonly struct KnownTestsCallbacks : ICallbacks
     {
         public void OnBeforeSend()
@@ -126,21 +219,65 @@ private readonly struct KnownTestsQuery
         [JsonProperty("configurations")]
         public readonly TestsConfigurations Configurations;
 
-        public KnownTestsQuery(string service, string environment, string repositoryUrl, TestsConfigurations configurations)
+        [JsonProperty("page_info")]
+        public readonly PageInfoRequest PageInfo;
+
+        public KnownTestsQuery(string service, string environment, string repositoryUrl, TestsConfigurations configurations, PageInfoRequest pageInfo)
         {
             Service = service;
             Environment = environment;
             RepositoryUrl = repositoryUrl;
             Configurations = configurations;
+            PageInfo = pageInfo;
+        }
+    }
+
+    private readonly struct PageInfoRequest
+    {
+        [JsonProperty("page_state")]
+        public readonly string? PageState;
+
+        public PageInfoRequest(string? pageState)
+        {
+            PageState = pageState;
         }
     }
 
+    private readonly struct PageInfoResponse
+    {
+        [JsonProperty("cursor")]
+        public readonly string? Cursor;
+
+        [JsonProperty("size")]
+        public readonly int Size;
+
+        [JsonProperty("has_next")]
+        public readonly bool HasNext;
+    }
+
+    /// <summary>
+    /// Internal response type for deserializing individual pages, which includes page_info.
+    /// </summary>
+    private readonly struct KnownTestsPageResponse
+    {
+        [JsonProperty("tests")]
+        public readonly KnownTestsResponse.KnownTestsModules? Tests;
+
+        [JsonProperty("page_info")]
+        public readonly PageInfoResponse? PageInfo;
+    }
+
     public readonly struct KnownTestsResponse
     {
         [JsonProperty("tests")]
         public readonly KnownTestsModules? Tests;
 
-        public sealed class KnownTestsSuites : Dictionary<string, string[]?>
+        public KnownTestsResponse(KnownTestsModules? tests)
+        {
+            Tests = tests;
+        }
+
+        public sealed class KnownTestsSuites : Dictionary<string, List<string>?>
         {
         }
 

tracer/test/Datadog.Trace.ClrProfiler.IntegrationTests/CI/TestingFrameworkEvpTest.cs

@@ -459,6 +459,7 @@ protected virtual async Task ExecuteTestAsync(string packageVersion, string evpV
             agent.Configuration.Endpoints = agent.Configuration.Endpoints.Where(e => !e.Contains(evpVersionToRemove)).ToArray();
 
             const string correlationId = "2e8a36bda770b683345957cc6c15baf9";
+            var knownTestsPageIndex = 0;
             agent.EventPlatformProxyPayloadReceived += (sender, e) =>
             {
                 if (e.Value.PathAndQuery.EndsWith("api/v2/libraries/tests/services/setting"))
@@ -469,7 +470,24 @@ protected virtual async Task ExecuteTestAsync(string packageVersion, string evpV
 
                 if (e.Value.PathAndQuery.EndsWith("api/v2/ci/libraries/tests"))
                 {
-                    e.Value.Response = string.IsNullOrEmpty(testScenario.MockData.TestsJson) ? new MockTracerResponse(string.Empty, 404) : new MockTracerResponse(testScenario.MockData.TestsJson, 200);
+                    if (testScenario.MockData.KnownTestsJsonPages is { Length: > 0 } pages)
+                    {
+                        // Serve paginated responses sequentially
+                        var idx = knownTestsPageIndex++;
+                        if (idx < pages.Length)
+                        {
+                            e.Value.Response = new MockTracerResponse(pages[idx], 200);
+                        }
+                        else
+                        {
+                            e.Value.Response = new MockTracerResponse(string.Empty, 404);
+                        }
+                    }
+                    else
+                    {
+                        e.Value.Response = string.IsNullOrEmpty(testScenario.MockData.TestsJson) ? new MockTracerResponse(string.Empty, 404) : new MockTracerResponse(testScenario.MockData.TestsJson, 200);
+                    }
+
                     return;
                 }
 
@@ -576,11 +594,27 @@ public readonly struct MockData
         public readonly string TestsJson;
         public readonly string TestManagementTestsJson;
 
+        /// <summary>
+        /// Optional paginated known tests responses. When non-null, the handler returns these
+        /// pages sequentially instead of <see cref="TestsJson"/>. Each entry must be a complete
+        /// JSON response including page_info with cursor/has_next.
+        /// </summary>
+        public readonly string[]? KnownTestsJsonPages;
+
         public MockData(string settingsJson, string testsJson, string testManagementTestsJson)
         {
             SettingsJson = settingsJson;
             TestsJson = testsJson;
             TestManagementTestsJson = testManagementTestsJson;
+            KnownTestsJsonPages = null;
+        }
+
+        public MockData(string settingsJson, string[] knownTestsJsonPages, string testManagementTestsJson)
+        {
+            SettingsJson = settingsJson;
+            TestsJson = string.Empty;
+            TestManagementTestsJson = testManagementTestsJson;
+            KnownTestsJsonPages = knownTestsJsonPages;
         }
 
         public override string ToString()

tracer/test/Datadog.Trace.ClrProfiler.IntegrationTests/CI/XUnitEvpTests.cs

@@ -118,6 +118,53 @@ public static IEnumerable<object[]> GetDataForEarlyFlakeDetection()
                 1,
                 115,
                 "efd_with_test_bypass");
+
+            // EFD with paginated known tests (TraitPassTest arrives on page 2)
+            yield return row.Concat(
+                new MockData(
+                    GetSettingsJson("true", "true", "false", "0"),
+                    [
+                        """
+                        {
+                            "data":{
+                                "id":"lNemDTwOV8U",
+                                "type":"ci_app_libraries_tests",
+                                "attributes":{
+                                    "tests":{},
+                                    "page_info":{
+                                        "cursor":"page-2-cursor",
+                                        "size":0,
+                                        "has_next":true
+                                    }
+                                }
+                            }
+                        }
+                        """,
+                        """
+                        {
+                            "data":{
+                                "id":"lNemDTwOV8U",
+                                "type":"ci_app_libraries_tests",
+                                "attributes":{
+                                    "tests":{
+                                        "Samples.XUnitTests":{
+                                            "Samples.XUnitTests.TestSuite":["TraitPassTest"]
+                                        }
+                                    },
+                                    "page_info":{
+                                        "cursor":"",
+                                        "size":1,
+                                        "has_next":false
+                                    }
+                                }
+                            }
+                        }
+                        """
+                    ],
+                    string.Empty),
+                1,
+                115,
+                "efd_with_test_bypass_paginated");
         }
     }