[SqlClient] Refactor Enrich test (#2714)

Author: alanwest

test/OpenTelemetry.Instrumentation.SqlClient.Tests/SqlClientTests.cs

@@ -70,19 +70,14 @@ public void SqlClient_NamedOptions()
     [MemberData(nameof(SqlTestData.SqlClientCallsAreCollectedSuccessfullyCases), MemberType = typeof(SqlTestData))]
     public void SqlClientCallsAreCollectedSuccessfully(
         string beforeCommand,
-        string afterCommand,
         CommandType commandType,
         string commandText,
         bool captureTextCommandContent,
-        bool shouldEnrich = true,
         bool emitOldAttributes = true,
         bool emitNewAttributes = false,
         bool tracingEnabled = true,
         bool metricsEnabled = true)
     {
-        using var sqlConnection = new SqlConnection(TestConnectionString);
-        using var sqlCommand = sqlConnection.CreateCommand();
-
         var activities = new List<Activity>();
         var metrics = new List<Metric>();
 
@@ -94,11 +89,6 @@ public void SqlClientCallsAreCollectedSuccessfully(
             (opt) =>
             {
                 opt.SetDbStatementForText = captureTextCommandContent;
-                if (shouldEnrich)
-                {
-                    opt.Enrich = ActivityEnrichment;
-                }
-
                 opt.EmitOldAttributes = emitOldAttributes;
                 opt.EmitNewAttributes = emitNewAttributes;
             });
@@ -118,33 +108,7 @@ public void SqlClientCallsAreCollectedSuccessfully(
 
         try
         {
-            var operationId = Guid.NewGuid();
-            sqlCommand.CommandType = commandType;
-#pragma warning disable CA2100
-            sqlCommand.CommandText = commandText;
-#pragma warning restore CA2100
-
-            var beforeExecuteEventData = new
-            {
-                OperationId = operationId,
-                Command = sqlCommand,
-                Timestamp = (long?)1000000L,
-            };
-
-            this.fakeSqlClientDiagnosticSource.Write(
-                beforeCommand,
-                beforeExecuteEventData);
-
-            var afterExecuteEventData = new
-            {
-                OperationId = operationId,
-                Command = sqlCommand,
-                Timestamp = 2000000L,
-            };
-
-            this.fakeSqlClientDiagnosticSource.Write(
-                afterCommand,
-                afterExecuteEventData);
+            this.ExecuteCommand(commandType, commandText, false, beforeCommand);
         }
         finally
         {
@@ -158,12 +122,12 @@ public void SqlClientCallsAreCollectedSuccessfully(
         {
             activity = Assert.Single(activities);
             VerifyActivityData(
-                sqlCommand.CommandType,
-                sqlCommand.CommandText,
+                commandType,
+                commandText,
                 captureTextCommandContent,
                 false,
                 false,
-                shouldEnrich,
+                false,
                 activity,
                 emitOldAttributes,
                 emitNewAttributes);
@@ -236,15 +200,10 @@ public void SqlClientAddsConnectionLevelAttributes(
     [MemberData(nameof(SqlTestData.SqlClientErrorsAreCollectedSuccessfullyCases), MemberType = typeof(SqlTestData))]
     public void SqlClientErrorsAreCollectedSuccessfully(
         string beforeCommand,
-        string errorCommand,
-        bool shouldEnrich = true,
         bool recordException = false,
         bool tracingEnabled = true,
         bool metricsEnabled = true)
     {
-        using var sqlConnection = new SqlConnection(TestConnectionString);
-        using var sqlCommand = sqlConnection.CreateCommand();
-
         var activities = new List<Activity>();
         var metrics = new List<Metric>();
         var traceProviderBuilder = Sdk.CreateTracerProviderBuilder();
@@ -254,10 +213,6 @@ public void SqlClientErrorsAreCollectedSuccessfully(
             traceProviderBuilder.AddSqlClientInstrumentation(options =>
             {
                 options.RecordException = recordException;
-                if (shouldEnrich)
-                {
-                    options.Enrich = ActivityEnrichment;
-                }
             })
             .AddInMemoryExporter(activities);
         }
@@ -277,32 +232,7 @@ public void SqlClientErrorsAreCollectedSuccessfully(
 
         try
         {
-            var operationId = Guid.NewGuid();
-            sqlCommand.CommandText = "SP_GetOrders";
-            sqlCommand.CommandType = CommandType.StoredProcedure;
-
-            var beforeExecuteEventData = new
-            {
-                OperationId = operationId,
-                Command = sqlCommand,
-                Timestamp = (long?)1000000L,
-            };
-
-            this.fakeSqlClientDiagnosticSource.Write(
-                beforeCommand,
-                beforeExecuteEventData);
-
-            var commandErrorEventData = new
-            {
-                OperationId = operationId,
-                Command = sqlCommand,
-                Exception = new Exception("Boom!"),
-                Timestamp = 2000000L,
-            };
-
-            this.fakeSqlClientDiagnosticSource.Write(
-                errorCommand,
-                commandErrorEventData);
+            this.ExecuteCommand(CommandType.StoredProcedure, "SP_GetOrders", true, beforeCommand);
         }
         finally
         {
@@ -316,12 +246,12 @@ public void SqlClientErrorsAreCollectedSuccessfully(
         {
             activity = Assert.Single(activities);
             VerifyActivityData(
-                sqlCommand.CommandType,
-                sqlCommand.CommandText,
+                CommandType.StoredProcedure,
+                "SP_GetOrders",
                 false,
                 true,
                 recordException,
-                shouldEnrich,
+                false,
                 activity);
         }
         else
@@ -358,6 +288,45 @@ public void MicrosoftDataStartsActivityWithExpectedAttributes(SqlClientTestCase
         this.RunSqlClientTestCase(testCase, SqlClientDiagnosticListener.SqlMicrosoftBeforeExecuteCommand, SqlClientDiagnosticListener.SqlMicrosoftAfterExecuteCommand);
     }
 
+    [Theory]
+    [InlineData(true, false)]
+    [InlineData(false, false)]
+    [InlineData(true, true)]
+    [InlineData(false, true)]
+    public void ShouldEnrichWhenEnabled(bool shouldEnrich, bool error)
+    {
+        var activities = new List<Activity>();
+
+        using var traceProvider = Sdk.CreateTracerProviderBuilder()
+            .AddSqlClientInstrumentation(
+            (opt) =>
+            {
+                if (shouldEnrich)
+                {
+                    opt.Enrich = ActivityEnrichment;
+                }
+            })
+            .AddInMemoryExporter(activities)
+            .Build();
+
+        this.ExecuteCommand(CommandType.Text, "SELECT * FROM Foo", error, SqlClientDiagnosticListener.SqlDataBeforeExecuteCommand);
+        this.ExecuteCommand(CommandType.Text, "SELECT * FROM Foo", error, SqlClientDiagnosticListener.SqlMicrosoftBeforeExecuteCommand);
+
+        Assert.Equal(2, activities.Count);
+        if (shouldEnrich)
+        {
+            Assert.Contains("enriched", activities[0].Tags.Select(x => x.Key));
+            Assert.Contains("enriched", activities[1].Tags.Select(x => x.Key));
+            Assert.Equal("yes", activities[0].Tags.FirstOrDefault(tag => tag.Key == "enriched").Value);
+            Assert.Equal("yes", activities[1].Tags.FirstOrDefault(tag => tag.Key == "enriched").Value);
+        }
+        else
+        {
+            Assert.DoesNotContain(activities[0].Tags, tag => tag.Key == "enriched");
+            Assert.DoesNotContain(activities[1].Tags, tag => tag.Key == "enriched");
+        }
+    }
+
     [Fact]
     public void ShouldCollectTelemetryWhenFilterEvaluatesToTrue()
     {
@@ -691,6 +660,65 @@ private Activity[] RunCommandWithFilter(
 
         return [.. activities];
     }
+
+    private void ExecuteCommand(CommandType commandType, string commandText, bool error, string beforeCommand)
+    {
+        var afterCommand = beforeCommand == SqlClientDiagnosticListener.SqlDataBeforeExecuteCommand
+            ? SqlClientDiagnosticListener.SqlDataAfterExecuteCommand
+            : SqlClientDiagnosticListener.SqlMicrosoftAfterExecuteCommand;
+
+        var errorCommand = beforeCommand == SqlClientDiagnosticListener.SqlDataBeforeExecuteCommand
+            ? SqlClientDiagnosticListener.SqlDataWriteCommandError
+            : SqlClientDiagnosticListener.SqlMicrosoftWriteCommandError;
+
+        using var sqlConnection = new SqlConnection(TestConnectionString);
+        using var sqlCommand = sqlConnection.CreateCommand();
+
+        var operationId = Guid.NewGuid();
+        sqlCommand.CommandType = commandType;
+#pragma warning disable CA2100
+        sqlCommand.CommandText = commandText;
+#pragma warning restore CA2100
+
+        var beforeExecuteEventData = new
+        {
+            OperationId = operationId,
+            Command = sqlCommand,
+            Timestamp = (long?)1000000L,
+        };
+
+        this.fakeSqlClientDiagnosticSource.Write(
+            beforeCommand,
+            beforeExecuteEventData);
+
+        if (error)
+        {
+            var commandErrorEventData = new
+            {
+                OperationId = operationId,
+                Command = sqlCommand,
+                Exception = new Exception("Boom!"),
+                Timestamp = 2000000L,
+            };
+
+            this.fakeSqlClientDiagnosticSource.Write(
+                errorCommand,
+                commandErrorEventData);
+        }
+        else
+        {
+            var afterExecuteEventData = new
+            {
+                OperationId = operationId,
+                Command = sqlCommand,
+                Timestamp = 2000000L,
+            };
+
+            this.fakeSqlClientDiagnosticSource.Write(
+                afterCommand,
+                afterExecuteEventData);
+        }
+    }
 #endif
 
     private class FakeSqlClientDiagnosticSource : IDisposable

test/OpenTelemetry.Instrumentation.SqlClient.Tests/SqlTestData.cs

@@ -17,26 +17,20 @@ public static IEnumerable<object[]> SqlClientCallsAreCollectedSuccessfullyCases(
         return from beforeCommand in new[] { SqlClientDiagnosticListener.SqlDataBeforeExecuteCommand, SqlClientDiagnosticListener.SqlMicrosoftBeforeExecuteCommand }
                from commandType in new[] { CommandType.StoredProcedure, CommandType.Text }
                from captureTextCommandContent in bools
-               from shouldEnrich in bools
                from emitOldAttributes in bools
                from emitNewAttributes in bools
                from tracingEnabled in bools
                from metricsEnabled in bools
                where emitOldAttributes && emitNewAttributes
-               let endCommand = beforeCommand == SqlClientDiagnosticListener.SqlDataBeforeExecuteCommand
-                   ? SqlClientDiagnosticListener.SqlDataAfterExecuteCommand
-                   : SqlClientDiagnosticListener.SqlMicrosoftAfterExecuteCommand
                let commandText = commandType == CommandType.Text
                    ? "select * from sys.databases"
                    : "SP_GetOrders"
                select new object[]
                {
                     beforeCommand,
-                    endCommand,
                     commandType,
                     commandText,
                     captureTextCommandContent,
-                    shouldEnrich,
                     emitOldAttributes,
                     emitNewAttributes,
                     tracingEnabled,
@@ -52,14 +46,9 @@ from shouldEnrich in bools
                from recordException in bools
                from tracingEnabled in bools
                from metricsEnabled in bools
-               let errorCommand = beforeCommand == SqlClientDiagnosticListener.SqlDataBeforeExecuteCommand
-                   ? SqlClientDiagnosticListener.SqlDataWriteCommandError
-                   : SqlClientDiagnosticListener.SqlMicrosoftWriteCommandError
                select new object[]
                {
                     beforeCommand,
-                    errorCommand,
-                    shouldEnrich,
                     recordException,
                     tracingEnabled,
                     metricsEnabled,