Refactor MockTracerAgent to allow sending custom responses for any en…

…dpoint

tracer/test/Datadog.Trace.ClrProfiler.IntegrationTests/AgentMalfunctionTests.cs

@@ -61,7 +61,22 @@ public async Task SubmitsTraces(AgentBehaviour behaviour, TestTransports transpo
 
             EnvironmentHelper.EnableTransport(transportType);
             using var agent = EnvironmentHelper.GetMockAgent();
-            agent.SetBehaviour(behaviour);
+            var customResponse = behaviour switch
+            {
+                AgentBehaviour.Return404 => new MockTracerResponse { StatusCode = 404 },
+                AgentBehaviour.Return500 => new MockTracerResponse { StatusCode = 500 },
+                AgentBehaviour.WrongAnswer => new MockTracerResponse("WRONG_ANSWER"),
+                AgentBehaviour.NoAnswer => new MockTracerResponse { SendResponse = false },
+                _ => null,
+            };
+
+            if (customResponse is { } cr)
+            {
+                // set everything except traces, but only these are actually used
+                agent.CustomResponses[MockTracerResponseType.Telemetry] = cr;
+                agent.CustomResponses[MockTracerResponseType.Info] = cr;
+                agent.CustomResponses[MockTracerResponseType.RemoteConfig] = cr;
+            }
 
             // The server implementation of named pipes is flaky so have 3 attempts
             var attemptsRemaining = 3;

tracer/test/Datadog.Trace.TestHelpers/MockTracerAgent.cs

@@ -36,8 +36,6 @@ public abstract class MockTracerAgent : IDisposable
     {
         private readonly CancellationTokenSource _cancellationTokenSource = new();
 
-        private AgentBehaviour behaviour = AgentBehaviour.Normal;
-
         protected MockTracerAgent(bool telemetryEnabled, TestTransports transport)
         {
             TelemetryEnabled = telemetryEnabled;
@@ -62,7 +60,7 @@ protected MockTracerAgent(bool telemetryEnabled, TestTransports transport)
 
         public bool TelemetryEnabled { get; }
 
-        public string RcmResponse { get; set; }
+        public Dictionary<MockTracerResponseType, MockTracerResponse> CustomResponses { get; } = new();
 
         /// <summary>
         /// Gets the filters used to filter out spans we don't want to look at for a test.
@@ -438,8 +436,6 @@ public virtual void Dispose()
             _cancellationTokenSource.Cancel();
         }
 
-        public void SetBehaviour(AgentBehaviour behaviour) => this.behaviour = behaviour;
-
         protected void IgnoreException(Action action)
         {
             try
@@ -474,74 +470,53 @@ protected virtual void OnMetricsReceived(string stats)
 
         private protected MockTracerResponse HandleHttpRequest(MockHttpParser.MockHttpRequest request)
         {
-            string response;
-            int statusCode;
-            bool sendResponse;
-            var isTraceCommand = false;
+            string response = null;
+            var responseType = MockTracerResponseType.Unknown;
 
             if (TelemetryEnabled && request.PathAndQuery.StartsWith("/" + TelemetryConstants.AgentTelemetryEndpoint))
             {
                 HandlePotentialTelemetryData(request);
-                response = "{}";
+                responseType = MockTracerResponseType.Telemetry;
             }
             else if (request.PathAndQuery.EndsWith("/info"))
             {
                 response = JsonConvert.SerializeObject(Configuration);
+                responseType = MockTracerResponseType.Info;
             }
             else if (request.PathAndQuery.StartsWith("/debugger/v1/input"))
             {
                 HandlePotentialDebuggerData(request);
-                response = "{}";
+                responseType = MockTracerResponseType.Debugger;
             }
             else if (request.PathAndQuery.StartsWith("/v0.6/stats"))
             {
                 HandlePotentialStatsData(request);
-                response = "{}";
+                responseType = MockTracerResponseType.Stats;
             }
             else if (request.PathAndQuery.StartsWith("/v0.7/config"))
             {
                 HandlePotentialRemoteConfig(request);
-                response = RcmResponse ?? "{}";
+                responseType = MockTracerResponseType.RemoteConfig;
             }
             else if (request.PathAndQuery.StartsWith("/v0.1/pipeline_stats"))
             {
                 HandlePotentialDataStreams(request);
-                response = "{}";
+                responseType = MockTracerResponseType.DataStreams;
             }
             else if (request.PathAndQuery.StartsWith("/evp_proxy/v2/"))
             {
                 HandleEvpProxyPayload(request);
-                response = "{}";
+                responseType = MockTracerResponseType.EvpProxy;
             }
             else
             {
                 HandlePotentialTraces(request);
-                response = "{}";
-                isTraceCommand = true;
-            }
-
-            if (isTraceCommand)
-            {
-                statusCode = 200;
-                sendResponse = true;
-            }
-            else
-            {
-                if (behaviour == AgentBehaviour.WrongAnswer)
-                {
-                    response = "WRONG_ANSWER";
-                }
-
-                sendResponse = behaviour != AgentBehaviour.NoAnswer;
-                statusCode = behaviour == AgentBehaviour.Return500 ? 500 : (behaviour == AgentBehaviour.Return404 ? 404 : 200);
+                responseType = MockTracerResponseType.Traces;
             }
 
-            return new MockTracerResponse()
-            {
-                Response = response,
-                SendResponse = sendResponse,
-                StatusCode = statusCode
-            };
+            return CustomResponses.TryGetValue(responseType, out var custom)
+                       ? custom // custom response, use that
+                       : new MockTracerResponse(response ?? "{}");
         }
 
         private void HandlePotentialTraces(MockHttpParser.MockHttpRequest request)

tracer/test/Datadog.Trace.TestHelpers/MockTracerResponse.cs

@@ -7,9 +7,24 @@ namespace Datadog.Trace.TestHelpers
 {
     public class MockTracerResponse
     {
+        public MockTracerResponse()
+        {
+        }
+
+        public MockTracerResponse(string response)
+        {
+            Response = response;
+        }
+
+        public MockTracerResponse(string response, int statusCode)
+        {
+            Response = response;
+            StatusCode = statusCode;
+        }
+
         public int StatusCode { get; set; } = 200;
 
-        public string Response { get; set; }
+        public string Response { get; set; } = "{}";
 
         public bool SendResponse { get; set; } = true;
     }

tracer/test/Datadog.Trace.TestHelpers/MockTracerResponseType.cs

@@ -0,0 +1,54 @@
+// <copyright file="MockTracerResponseType.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>
+
+namespace Datadog.Trace.TestHelpers;
+
+public enum MockTracerResponseType
+{
+    /// <summary>
+    /// Any request which doesn't match a known endpoint
+    /// </summary>
+    Unknown,
+
+    /// <summary>
+    /// The trace endpoint
+    /// </summary>
+    Traces,
+
+    /// <summary>
+    /// The Telemetry endpoint
+    /// </summary>
+    Telemetry,
+
+    /// <summary>
+    /// The discovery endpoint
+    /// </summary>
+    Info,
+
+    /// <summary>
+    /// The dynamic configuration endpoint
+    /// </summary>
+    Debugger,
+
+    /// <summary>
+    /// The trace stats endpoint
+    /// </summary>
+    Stats,
+
+    /// <summary>
+    /// The remote configuration endpoint
+    /// </summary>
+    RemoteConfig,
+
+    /// <summary>
+    /// The Data streams endpoint
+    /// </summary>
+    DataStreams,
+
+    /// <summary>
+    /// The CI Visibility EVP proxy endpoint
+    /// </summary>
+    EvpProxy,
+}

tracer/test/Datadog.Trace.TestHelpers/RemoteConfigTestHelper.cs

@@ -23,15 +23,15 @@ public static class RemoteConfigTestHelper
         internal static GetRcmResponse SetupRcm(this MockTracerAgent agent, ITestOutputHelper output, IEnumerable<(object Config, string ProductName, string Id)> configurations)
         {
             var response = BuildRcmResponse(configurations.Select(c => (JsonConvert.SerializeObject(c.Config), c.ProductName, c.Id)));
-            agent.RcmResponse = JsonConvert.SerializeObject(response);
+            agent.CustomResponses[MockTracerResponseType.RemoteConfig] = new(JsonConvert.SerializeObject(response));
             output.WriteLine($"{DateTime.UtcNow}: Using RCM response: {response}");
             return response;
         }
 
         internal static async Task<GetRcmRequest> SetupRcmAndWait(this MockTracerAgent agent, ITestOutputHelper output, IEnumerable<(object Config, string ProductName, string Id)> configurations, int timeoutInMilliseconds = WaitForAcknowledgmentTimeout)
         {
             var response = BuildRcmResponse(configurations.Select(c => (JsonConvert.SerializeObject(c.Config), c.ProductName, c.Id)));
-            agent.RcmResponse = JsonConvert.SerializeObject(response);
+            agent.CustomResponses[MockTracerResponseType.RemoteConfig] = new(JsonConvert.SerializeObject(response));
             output.WriteLine($"{DateTime.UtcNow}: Using RCM response: {response} with custom opaque state {response.Targets.Signed.Custom.OpaqueBackendState}");
             var res = await agent.WaitRcmRequestAndReturnMatchingRequest(response, timeoutInMilliseconds: timeoutInMilliseconds);
             return res;