Refactor tests of StringBuffer substring (#8000)

DataDog · Nov 22, 2024 · 9078a8c · 9078a8c
1 parent bdfcc9b
commit 9078a8c
Show file tree

Hide file tree

Showing 3 changed files with 48 additions and 14 deletions.
diff --git a/.../src/test/groovy/datadog/trace/instrumentation/java/lang/StringBuilderCallSiteTest.groovy b/.../src/test/groovy/datadog/trace/instrumentation/java/lang/StringBuilderCallSiteTest.groovy
@@ -190,6 +190,23 @@ class StringBuilderCallSiteTest extends AgentTestRunner {
     where:
     param         | beginIndex | expected
     sb('012345')  | 1          | '12345'
+  }
+
+  def 'test string buffer substring call site'() {
+    setup:
+    final iastModule = Mock(StringModule)
+    InstrumentationBridge.registerIastModule(iastModule)
+
+    when:
+    final result = TestStringBufferSuite.substring(param, beginIndex)
+
+    then:
+    result == expected
+    1 * iastModule.onStringSubSequence(param, beginIndex, param.length(), expected)
+    0 * _
+
+    where:
+    param         | beginIndex | expected
     sbf('012345') | 1          | '12345'
   }
 
@@ -209,6 +226,23 @@ class StringBuilderCallSiteTest extends AgentTestRunner {
     where:
     param         | beginIndex | endIndex | expected
     sb('012345')  | 1          | 5        | '1234'
+  }
+
+  def 'test string buffer substring with endIndex call site'() {
+    setup:
+    final iastModule = Mock(StringModule)
+    InstrumentationBridge.registerIastModule(iastModule)
+
+    when:
+    final result = TestStringBufferSuite.substring(param, beginIndex, endIndex)
+
+    then:
+    result == expected
+    1 * iastModule.onStringSubSequence(param, beginIndex, endIndex, expected)
+    0 * _
+
+    where:
+    param         | beginIndex | endIndex | expected
     sbf('012345') | 1          | 5        | '1234'
   }
 

diff --git a/dd-java-agent/instrumentation/java-lang/src/test/java/foo/bar/TestStringBufferSuite.java b/dd-java-agent/instrumentation/java-lang/src/test/java/foo/bar/TestStringBufferSuite.java
@@ -51,4 +51,18 @@ public String toString(final StringBuffer buffer) {
     LOGGER.debug("After string buffer toString {}", result);
     return result;
   }
+
+  public static String substring(StringBuffer self, int beginIndex, int endIndex) {
+    LOGGER.debug("Before string buffer substring {} from {} to {}", self, beginIndex, endIndex);
+    final String result = self.substring(beginIndex, endIndex);
+    LOGGER.debug("After string buffer substring {}", result);
+    return result;
+  }
+
+  public static String substring(StringBuffer self, int beginIndex) {
+    LOGGER.debug("Before string buffer substring {} from {}", self, beginIndex);
+    final String result = self.substring(beginIndex);
+    LOGGER.debug("After string buffer substring {}", result);
+    return result;
+  }
 }
diff --git a/dd-java-agent/instrumentation/java-lang/src/test/java/foo/bar/TestStringBuilderSuite.java b/dd-java-agent/instrumentation/java-lang/src/test/java/foo/bar/TestStringBuilderSuite.java
@@ -90,18 +90,4 @@ public static String substring(StringBuilder self, int beginIndex) {
     LOGGER.debug("After string builder substring {}", result);
     return result;
   }
-
-  public static String substring(StringBuffer self, int beginIndex, int endIndex) {
-    LOGGER.debug("Before string buffer substring {} from {} to {}", self, beginIndex, endIndex);
-    final String result = self.substring(beginIndex, endIndex);
-    LOGGER.debug("After string buffer substring {}", result);
-    return result;
-  }
-
-  public static String substring(StringBuffer self, int beginIndex) {
-    LOGGER.debug("Before string buffer substring {} from {}", self, beginIndex);
-    final String result = self.substring(beginIndex);
-    LOGGER.debug("After string buffer substring {}", result);
-    return result;
-  }
 }