feat: pass the force record flag to storage service (#428)

pangdayuan1 · web-flow · commit 11ee390fc913 · 2024-03-22T20:09:06.000+08:00
diff --git a/arex-instrumentation-api/src/main/java/io/arex/inst/runtime/context/ArexContext.java b/arex-instrumentation-api/src/main/java/io/arex/inst/runtime/context/ArexContext.java
@@ -91,6 +91,10 @@ public void setExcludeMockTemplate(Map<String, Set<String>> excludeMockTemplate)
     }
 
     public void setAttachment(String key, Object value) {
+        if (value == null) {
+            return;
+        }
+
         if (attachments == null) {
             attachments = new HashMap<>();
         }
diff --git a/arex-instrumentation-api/src/test/java/io/arex/inst/runtime/context/ArexContextTest.java b/arex-instrumentation-api/src/test/java/io/arex/inst/runtime/context/ArexContextTest.java
@@ -11,4 +11,11 @@ void calculateSequence() {
         ArexContext arexContext = ArexContext.of("mock");
         assertEquals(0, arexContext.calculateSequence());
     }
+
+    @Test
+    void setAttachment() {
+        ArexContext arexContext = ArexContext.of("mock");
+        arexContext.setAttachment("testKey", null);
+        assertNull(arexContext.getAttachment("testKey"));
+    }
 }
diff --git a/arex-instrumentation/dubbo/arex-dubbo-alibaba/src/main/java/io/arex/inst/dubbo/alibaba/DubboProviderExtractor.java b/arex-instrumentation/dubbo/arex-dubbo-alibaba/src/main/java/io/arex/inst/dubbo/alibaba/DubboProviderExtractor.java
@@ -27,6 +27,7 @@ public static void onServiceEnter(Invoker<?> invoker, Invocation invocation) {
         String excludeMockTemplate = adapter.getExcludeMockTemplate();
         RequestHandlerManager.preHandle(invocation.getAttachments(), MockCategoryType.DUBBO_PROVIDER.getName());
         CaseEventDispatcher.onEvent(CaseEvent.ofCreateEvent(EventSource.of(caseId, excludeMockTemplate)));
+        ContextManager.currentContext().setAttachment(ArexConstants.FORCE_RECORD, adapter.forceRecord());
         RequestHandlerManager.handleAfterCreateContext(invocation.getAttachments(), MockCategoryType.DUBBO_PROVIDER.getName());
         invocation.getAttachments().put(ArexConstants.ORIGINAL_REQUEST, Serializer.serialize(invocation.getArguments()));
     }
diff --git a/arex-instrumentation/dubbo/arex-dubbo-alibaba/src/test/java/io/arex/inst/dubbo/alibaba/DubboProviderExtractorTest.java b/arex-instrumentation/dubbo/arex-dubbo-alibaba/src/test/java/io/arex/inst/dubbo/alibaba/DubboProviderExtractorTest.java
@@ -42,6 +42,7 @@ static void setUp() {
         Mockito.mockStatic(RecordLimiter.class);
         Mockito.mockStatic(IgnoreUtils.class);
         Mockito.mockStatic(Config.class);
+        Mockito.when(ContextManager.currentContext()).thenReturn(ArexContext.of("mock"));
         Mockito.when(Config.get()).thenReturn(Mockito.mock(Config.class));
     }
 
diff --git a/arex-instrumentation/dubbo/arex-dubbo-apache-v2/src/main/java/io/arex/inst/dubbo/apache/v2/DubboProviderExtractor.java b/arex-instrumentation/dubbo/arex-dubbo-apache-v2/src/main/java/io/arex/inst/dubbo/apache/v2/DubboProviderExtractor.java
@@ -30,6 +30,7 @@ public static void onServiceEnter(Invoker<?> invoker, Invocation invocation) {
         String excludeMockTemplate = adapter.getExcludeMockTemplate();
         RequestHandlerManager.preHandle(invocation.getAttachments(), MockCategoryType.DUBBO_PROVIDER.getName());
         CaseEventDispatcher.onEvent(CaseEvent.ofCreateEvent(EventSource.of(caseId, excludeMockTemplate)));
+        ContextManager.currentContext().setAttachment(ArexConstants.FORCE_RECORD, adapter.forceRecord());
         RequestHandlerManager.handleAfterCreateContext(invocation.getAttachments(), MockCategoryType.DUBBO_PROVIDER.getName());
         invocation.setAttachment(ArexConstants.ORIGINAL_REQUEST, Serializer.serialize(invocation.getArguments()));
     }
diff --git a/arex-instrumentation/dubbo/arex-dubbo-apache-v2/src/test/java/io/arex/inst/dubbo/apache/v2/DubboProviderExtractorTest.java b/arex-instrumentation/dubbo/arex-dubbo-apache-v2/src/test/java/io/arex/inst/dubbo/apache/v2/DubboProviderExtractorTest.java
@@ -42,6 +42,7 @@ static void setUp() {
         Mockito.mockStatic(IgnoreUtils.class);
         Mockito.mockStatic(Config.class);
         Mockito.when(Config.get()).thenReturn(Mockito.mock(Config.class));
+        Mockito.when(ContextManager.currentContext()).thenReturn(ArexContext.of("mock"));
     }
 
     @AfterAll
diff --git a/arex-instrumentation/dubbo/arex-dubbo-apache-v3/src/main/java/io/arex/inst/dubbo/apache/v3/DubboProviderExtractor.java b/arex-instrumentation/dubbo/arex-dubbo-apache-v3/src/main/java/io/arex/inst/dubbo/apache/v3/DubboProviderExtractor.java
@@ -27,6 +27,7 @@ public static void onServiceEnter(Invoker<?> invoker, Invocation invocation) {
         String excludeMockTemplate = adapter.getExcludeMockTemplate();
         RequestHandlerManager.preHandle(invocation.getAttachments(), MockCategoryType.DUBBO_PROVIDER.getName());
         CaseEventDispatcher.onEvent(CaseEvent.ofCreateEvent(EventSource.of(caseId, excludeMockTemplate)));
+        ContextManager.currentContext().setAttachment(ArexConstants.FORCE_RECORD, adapter.forceRecord());
         RequestHandlerManager.handleAfterCreateContext(invocation.getAttachments(), MockCategoryType.DUBBO_PROVIDER.getName());
         invocation.getAttributes().put(ArexConstants.ORIGINAL_REQUEST, Serializer.serialize(invocation.getArguments()));
     }
diff --git a/arex-instrumentation/dubbo/arex-dubbo-apache-v3/src/test/java/io/arex/inst/dubbo/apache/v3/DubboProviderExtractorTest.java b/arex-instrumentation/dubbo/arex-dubbo-apache-v3/src/test/java/io/arex/inst/dubbo/apache/v3/DubboProviderExtractorTest.java
@@ -42,6 +42,7 @@ static void setUp() {
         Mockito.mockStatic(IgnoreUtils.class);
         Mockito.mockStatic(Config.class);
         Mockito.when(Config.get()).thenReturn(Mockito.mock(Config.class));
+        Mockito.when(ContextManager.currentContext()).thenReturn(ArexContext.of("mock"));
     }
 
     @AfterAll
diff --git a/arex-instrumentation/netty/arex-netty-v3/src/main/java/io/arex/inst/netty/v3/server/RequestTracingHandler.java b/arex-instrumentation/netty/arex-netty-v3/src/main/java/io/arex/inst/netty/v3/server/RequestTracingHandler.java
@@ -34,6 +34,7 @@ public void messageReceived(ChannelHandlerContext ctx, MessageEvent event) throw
 
                 String excludeMockTemplate = NettyHelper.getHeader(request, ArexConstants.HEADER_EXCLUDE_MOCK);
                 CaseEventDispatcher.onEvent(CaseEvent.ofCreateEvent(EventSource.of(caseId, excludeMockTemplate)));
+                ContextManager.currentContext().setAttachment(ArexConstants.FORCE_RECORD, NettyHelper.getHeader(request, ArexConstants.FORCE_RECORD));
                 if (ContextManager.needRecordOrReplay()) {
                     Mocker mocker = MockUtils.createNettyProvider(request.getUri());
                     Mocker.Target target = mocker.getTargetRequest();
diff --git a/arex-instrumentation/netty/arex-netty-v3/src/test/java/io/arex/inst/netty/v3/server/RequestTracingHandlerTest.java b/arex-instrumentation/netty/arex-netty-v3/src/test/java/io/arex/inst/netty/v3/server/RequestTracingHandlerTest.java
@@ -52,6 +52,7 @@ static void setUp() {
         Mockito.mockStatic(Config.class);
         Mockito.when(Config.get()).thenReturn(Mockito.mock(Config.class));
         Mockito.mockStatic(MockUtils.class);
+        Mockito.when(ContextManager.currentContext()).thenReturn(ArexContext.of("mock"));
     }
 
     @AfterAll
diff --git a/arex-instrumentation/netty/arex-netty-v4/src/main/java/io/arex/inst/netty/v4/server/RequestTracingHandler.java b/arex-instrumentation/netty/arex-netty-v4/src/main/java/io/arex/inst/netty/v4/server/RequestTracingHandler.java
@@ -36,6 +36,7 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception
 
                 String excludeMockTemplate = request.headers().get(ArexConstants.HEADER_EXCLUDE_MOCK);
                 CaseEventDispatcher.onEvent(CaseEvent.ofCreateEvent(EventSource.of(caseId, excludeMockTemplate)));
+                ContextManager.currentContext().setAttachment(ArexConstants.FORCE_RECORD, request.headers().get(ArexConstants.FORCE_RECORD));
                 if (ContextManager.needRecordOrReplay()) {
                     Mocker mocker = MockUtils.createNettyProvider(request.getUri());
                     Mocker.Target target = mocker.getTargetRequest();
diff --git a/arex-instrumentation/netty/arex-netty-v4/src/test/java/io/arex/inst/netty/v4/server/RequestTracingHandlerTest.java b/arex-instrumentation/netty/arex-netty-v4/src/test/java/io/arex/inst/netty/v4/server/RequestTracingHandlerTest.java
@@ -4,6 +4,7 @@
 import io.arex.agent.bootstrap.model.Mocker.Target;
 import io.arex.agent.bootstrap.util.Assert;
 import io.arex.inst.runtime.config.Config;
+import io.arex.inst.runtime.context.ArexContext;
 import io.arex.inst.runtime.context.ContextManager;
 import io.arex.inst.runtime.context.RecordLimiter;
 import io.arex.inst.runtime.listener.CaseEventDispatcher;
@@ -57,6 +58,7 @@ static void setUp() {
         Mockito.when(Config.get()).thenReturn(Mockito.mock(Config.class));
         mockCaseEvent = Mockito.mockStatic(CaseEventDispatcher.class);
         Mockito.mockStatic(MockUtils.class);
+        Mockito.when(ContextManager.currentContext()).thenReturn(ArexContext.of("mock"));
     }
 
     @AfterAll
diff --git a/arex-instrumentation/servlet/arex-httpservlet/src/main/java/io/arex/inst/httpservlet/ServletAdviceHelper.java b/arex-instrumentation/servlet/arex-httpservlet/src/main/java/io/arex/inst/httpservlet/ServletAdviceHelper.java
@@ -119,6 +119,8 @@ public static <TRequest, TResponse> Pair<TRequest, TResponse> onServiceEnter(
             String caseId = adapter.getRequestHeader(httpServletRequest, ArexConstants.RECORD_ID);
             String excludeMockTemplate = adapter.getRequestHeader(httpServletRequest, ArexConstants.HEADER_EXCLUDE_MOCK);
             CaseEventDispatcher.onEvent(CaseEvent.ofCreateEvent(EventSource.of(caseId, excludeMockTemplate)));
+            ContextManager.currentContext().setAttachment(ArexConstants.FORCE_RECORD,
+                adapter.getRequestHeader(httpServletRequest, ArexConstants.FORCE_RECORD));
             RequestHandlerManager.handleAfterCreateContext(httpServletRequest, MockCategoryType.SERVLET.getName());
         }
 
diff --git a/arex-instrumentation/servlet/arex-httpservlet/src/test/java/io/arex/inst/httpservlet/ServletAdviceHelperTest.java b/arex-instrumentation/servlet/arex-httpservlet/src/test/java/io/arex/inst/httpservlet/ServletAdviceHelperTest.java
@@ -100,13 +100,17 @@ static Stream<Arguments> onServiceEnterCase() {
         Runnable shouldSkip1 = () -> {
             Mockito.when(adapter.getMethod(any())).thenReturn("POST");
             Mockito.when(adapter.getRequestHeader(any(), eq(ArexConstants.RECORD_ID))).thenReturn("mock");
+            Mockito.when(ContextManager.currentContext()).thenReturn(ArexContext.of("mock"));
+            Mockito.when(adapter.getRequestHeader(any(), eq(ArexConstants.FORCE_RECORD))).thenReturn("true");
         };
         Runnable getRedirectRecordId1 = () -> {
             Mockito.when(adapter.getParameterFromQueryString(any(), eq(ArexConstants.RECORD_ID))).thenReturn(null);
+            Mockito.when(adapter.getRequestHeader(any(), eq(ArexConstants.FORCE_RECORD))).thenReturn("true");
         };
         Runnable getRedirectRecordId2 = () -> {
             Mockito.when(adapter.getParameterFromQueryString(any(), eq(ArexConstants.RECORD_ID))).thenReturn("mock-redirectRecordId");
             Mockito.when(adapter.getRequestHeader(any(), eq("referer"))).thenReturn(null);
+            Mockito.when(adapter.getRequestHeader(any(), eq("FORCE_RECORD"))).thenReturn(null);
         };
         Runnable getRedirectRecordId3 = () -> {
             Mockito.when(adapter.getParameterFromQueryString(any(), eq(ArexConstants.RECORD_ID))).thenReturn("mock-redirectRecordId");

Original file line number	Diff line number	Diff line change
`@@ -91,6 +91,10 @@ public void setExcludeMockTemplate(Map<String, Set<String>> excludeMockTemplate)`
`91`	`91`	`}`
`92`	`92`
`93`	`93`	`public void setAttachment(String key, Object value) {`
	`94`	`+ if (value == null) {`
	`95`	`+ return;`
	`96`	`+ }`
	`97`	`+`
`94`	`98`	`if (attachments == null) {`
`95`	`99`	`attachments = new HashMap<>();`
`96`	`100`	`}`
Original file line number	Diff line number	Diff line change
`@@ -11,4 +11,11 @@ void calculateSequence() {`
`11`	`11`	`ArexContext arexContext = ArexContext.of("mock");`
`12`	`12`	`assertEquals(0, arexContext.calculateSequence());`
`13`	`13`	`}`
	`14`	`+`
	`15`	`+ @Test`
	`16`	`+ void setAttachment() {`
	`17`	`+ ArexContext arexContext = ArexContext.of("mock");`
	`18`	`+ arexContext.setAttachment("testKey", null);`
	`19`	`+ assertNull(arexContext.getAttachment("testKey"));`
	`20`	`+ }`
`14`	`21`	`}`
Original file line number	Diff line number	Diff line change
`@@ -27,6 +27,7 @@ public static void onServiceEnter(Invoker<?> invoker, Invocation invocation) {`
`27`	`27`	`String excludeMockTemplate = adapter.getExcludeMockTemplate();`
`28`	`28`	`RequestHandlerManager.preHandle(invocation.getAttachments(), MockCategoryType.DUBBO_PROVIDER.getName());`
`29`	`29`	`CaseEventDispatcher.onEvent(CaseEvent.ofCreateEvent(EventSource.of(caseId, excludeMockTemplate)));`
	`30`	`+ ContextManager.currentContext().setAttachment(ArexConstants.FORCE_RECORD, adapter.forceRecord());`
`30`	`31`	`RequestHandlerManager.handleAfterCreateContext(invocation.getAttachments(), MockCategoryType.DUBBO_PROVIDER.getName());`
`31`	`32`	`invocation.getAttachments().put(ArexConstants.ORIGINAL_REQUEST, Serializer.serialize(invocation.getArguments()));`
`32`	`33`	`}`
Original file line number	Diff line number	Diff line change
`@@ -42,6 +42,7 @@ static void setUp() {`
`42`	`42`	`Mockito.mockStatic(RecordLimiter.class);`
`43`	`43`	`Mockito.mockStatic(IgnoreUtils.class);`
`44`	`44`	`Mockito.mockStatic(Config.class);`
	`45`	`+ Mockito.when(ContextManager.currentContext()).thenReturn(ArexContext.of("mock"));`
`45`	`46`	`Mockito.when(Config.get()).thenReturn(Mockito.mock(Config.class));`
`46`	`47`	`}`
`47`	`48`
Original file line number	Diff line number	Diff line change
`@@ -30,6 +30,7 @@ public static void onServiceEnter(Invoker<?> invoker, Invocation invocation) {`
`30`	`30`	`String excludeMockTemplate = adapter.getExcludeMockTemplate();`
`31`	`31`	`RequestHandlerManager.preHandle(invocation.getAttachments(), MockCategoryType.DUBBO_PROVIDER.getName());`
`32`	`32`	`CaseEventDispatcher.onEvent(CaseEvent.ofCreateEvent(EventSource.of(caseId, excludeMockTemplate)));`
	`33`	`+ ContextManager.currentContext().setAttachment(ArexConstants.FORCE_RECORD, adapter.forceRecord());`
`33`	`34`	`RequestHandlerManager.handleAfterCreateContext(invocation.getAttachments(), MockCategoryType.DUBBO_PROVIDER.getName());`
`34`	`35`	`invocation.setAttachment(ArexConstants.ORIGINAL_REQUEST, Serializer.serialize(invocation.getArguments()));`
`35`	`36`	`}`
Original file line number	Diff line number	Diff line change
`@@ -52,6 +52,7 @@ static void setUp() {`
`52`	`52`	`Mockito.mockStatic(Config.class);`
`53`	`53`	`Mockito.when(Config.get()).thenReturn(Mockito.mock(Config.class));`
`54`	`54`	`Mockito.mockStatic(MockUtils.class);`
	`55`	`+ Mockito.when(ContextManager.currentContext()).thenReturn(ArexContext.of("mock"));`
`55`	`56`	`}`
`56`	`57`
`57`	`58`	`@AfterAll`