Improved unit test coverage and minor refactor.

newrelic · Jan 21, 2025 · b9b5615 · b9b5615
1 parent 3f8fd5f
commit b9b5615
Show file tree

Hide file tree

Showing 4 changed files with 245 additions and 78 deletions.
diff --git a/...t/NewRelic/Agent/Extensions/NewRelic.Agent.Extensions/Helpers/SearchRequestWrapperBase.cs b/...t/NewRelic/Agent/Extensions/NewRelic.Agent.Extensions/Helpers/SearchRequestWrapperBase.cs
@@ -22,29 +22,15 @@ public abstract class SearchRequestWrapperBase
         private Func<object, object> _exceptionGetter;
         private Func<object, Uri> _uriGetter;
 
-        public ConcurrentDictionary<Type, Func<object, object>> GetRequestResponseFromGeneric = new ConcurrentDictionary<Type, Func<object, object>>();
+        protected ConcurrentDictionary<Type, Func<object, object>> GetRequestResponseFromGeneric = new ConcurrentDictionary<Type, Func<object, object>>();
 
         public abstract DatastoreVendor Vendor { get; } 
 
-        public abstract int RequestParamsIndex { get; }
-
-        public abstract int RequestParamsIndexAsync { get; }
-
-        public bool IsTransactionRequired => true;
-
-        public ISegment BuildSegment(bool isAsync, InstrumentedMethodCall instrumentedMethodCall, ITransaction transaction)
+        protected ISegment BuildSegment(int requestParamsIndex, InstrumentedMethodCall instrumentedMethodCall, ITransaction transaction)
         {
-            var indexOfRequestParams = RequestParamsIndex;
-
-            if (isAsync)
-            {
-                transaction.AttachToAsync();
-                indexOfRequestParams = RequestParamsIndexAsync;
-            }
-
             var path = (string)instrumentedMethodCall.MethodCall.MethodArguments[1];
             var request = instrumentedMethodCall.MethodCall.MethodArguments[0].ToString();
-            var requestParams = instrumentedMethodCall.MethodCall.MethodArguments[indexOfRequestParams];
+            var requestParams = instrumentedMethodCall.MethodCall.MethodArguments[requestParamsIndex];
             var splitPath = path.Trim('/').Split('/');
 
             var operation = (requestParams == null) ? GetOperationFromPath(request, splitPath) : GetOperationFromRequestParams(requestParams);
@@ -63,7 +49,7 @@ public ISegment BuildSegment(bool isAsync, InstrumentedMethodCall instrumentedMe
             return segment;
         }
 
-        public void TryProcessResponse(IAgent agent, ITransaction transaction, object response, ISegment segment)
+        protected void TryProcessResponse(IAgent agent, ITransaction transaction, object response, ISegment segment)
         {
             try
             {
@@ -280,7 +266,7 @@ private static void SetUriOnDatastoreSegment(ISegment segment, Uri uri)
             segmentExperimentalApi.SetSegmentData(data);
         }
 
-        public static bool ValidTaskResponse(Task response)
+        protected static bool ValidTaskResponse(Task response)
         {
             return response?.Status == TaskStatus.RanToCompletion;
         }

diff --git a/src/Agent/NewRelic/Agent/Extensions/Providers/Wrapper/Elasticsearch/RequestWrapper.cs b/src/Agent/NewRelic/Agent/Extensions/Providers/Wrapper/Elasticsearch/RequestWrapper.cs
@@ -16,12 +16,12 @@ namespace NewRelic.Providers.Wrapper.Elasticsearch
     public class RequestWrapper : SearchRequestWrapperBase, IWrapper
     {
         private const string WrapperName = "ElasticsearchRequestWrapper";
+        private const int RequestParamsIndex = 3;
+        private const int RequestParamsIndexAsync = 4;
 
         public override DatastoreVendor Vendor => DatastoreVendor.Elasticsearch;
 
-        public override int RequestParamsIndex => 3;
-
-        public override int RequestParamsIndexAsync => 4;
+        public bool IsTransactionRequired => true;
 
         public CanWrapResponse CanWrap(InstrumentedMethodInfo methodInfo)
         {
@@ -33,7 +33,15 @@ public AfterWrappedMethodDelegate BeforeWrappedMethod(InstrumentedMethodCall ins
             var isAsync = instrumentedMethodCall.IsAsync ||
                           instrumentedMethodCall.InstrumentedMethodInfo.Method.MethodName.EndsWith("RequestAsync");
 
-            var segment = BuildSegment(isAsync, instrumentedMethodCall, transaction);
+            var indexOfRequestParams = RequestParamsIndex;
+
+            if (isAsync)
+            {
+                transaction.AttachToAsync();
+                indexOfRequestParams = RequestParamsIndexAsync;
+            }
+
+            var segment = BuildSegment(indexOfRequestParams, instrumentedMethodCall, transaction);
 
             if (isAsync)
             {

diff --git a/src/Agent/NewRelic/Agent/Extensions/Providers/Wrapper/OpenSearch/RequestWrapper.cs b/src/Agent/NewRelic/Agent/Extensions/Providers/Wrapper/OpenSearch/RequestWrapper.cs
@@ -19,12 +19,12 @@ namespace NewRelic.Providers.Wrapper.OpenSearch
     public class RequestWrapper : SearchRequestWrapperBase, IWrapper
     {
         private const string WrapperName = "OpenSearchRequestWrapper";
+        private const int RequestParamsIndex = 3;
+        private const int RequestParamsIndexAsync = 4;
 
         public override DatastoreVendor Vendor => DatastoreVendor.OpenSearch;
 
-        public override int RequestParamsIndex => 3;
-
-        public override int RequestParamsIndexAsync => 4;
+        public bool IsTransactionRequired => true;
 
         public CanWrapResponse CanWrap(InstrumentedMethodInfo methodInfo)
         {
@@ -37,7 +37,15 @@ public AfterWrappedMethodDelegate BeforeWrappedMethod(InstrumentedMethodCall ins
             var isAsync = instrumentedMethodCall.IsAsync ||
                           instrumentedMethodCall.InstrumentedMethodInfo.Method.MethodName.EndsWith("RequestAsync");
 
-            var segment = BuildSegment(isAsync, instrumentedMethodCall, transaction);
+            var indexOfRequestParams = RequestParamsIndex;
+
+            if (isAsync)
+            {
+                transaction.AttachToAsync();
+                indexOfRequestParams = RequestParamsIndexAsync;
+            }
+
+            var segment = BuildSegment(indexOfRequestParams, instrumentedMethodCall, transaction);
 
             if (isAsync)
             {