Config and system properties preservation in ddspec

randomanderson · randomanderson · commit eace412422de · 2020-10-23T12:36:40.000-04:00
diff --git a/utils/test-utils/src/main/groovy/datadog/trace/test/util/DDSpecification.groovy b/utils/test-utils/src/main/groovy/datadog/trace/test/util/DDSpecification.groovy
@@ -1,11 +1,19 @@
 package datadog.trace.test.util
 
+
 import net.bytebuddy.agent.ByteBuddyAgent
 import net.bytebuddy.agent.builder.AgentBuilder
 import net.bytebuddy.dynamic.ClassFileLocator
 import net.bytebuddy.dynamic.Transformer
+import org.junit.Rule
+import org.junit.contrib.java.lang.system.EnvironmentVariables
+import org.junit.contrib.java.lang.system.RestoreSystemProperties
 import spock.lang.Specification
 
+import java.lang.reflect.Constructor
+import java.lang.reflect.Field
+import java.lang.reflect.Modifier
+
 import static net.bytebuddy.description.modifier.FieldManifestation.VOLATILE
 import static net.bytebuddy.description.modifier.Ownership.STATIC
 import static net.bytebuddy.description.modifier.Visibility.PUBLIC
@@ -15,13 +23,22 @@ import static net.bytebuddy.matcher.ElementMatchers.none
 abstract class DDSpecification extends Specification {
   private static final String CONFIG = "datadog.trace.api.Config"
 
+  private static Field CONFIG_INSTANCE_FIELD
+  private static Constructor CONFIG_CONSTRUCTOR
+
   static {
     makeConfigInstanceModifiable()
   }
 
   // Keep track of config instance already made modifiable
   private static isConfigInstanceModifiable = false
 
+  @Rule
+  public final EnvironmentVariables environmentVariables = new EnvironmentVariables()
+
+  @Rule
+  protected final RestoreSystemProperties restoreSystemProperties = new RestoreSystemProperties()
+
   static void makeConfigInstanceModifiable() {
     if (isConfigInstanceModifiable) {
       return
@@ -43,6 +60,46 @@ abstract class DDSpecification extends Specification {
           .transform(Transformer.ForField.withModifiers(PUBLIC, STATIC, VOLATILE))
       }
       .installOn(instrumentation)
+
+    Class configClass = Class.forName(CONFIG)
+    CONFIG_INSTANCE_FIELD = configClass.getDeclaredField("INSTANCE")
+    CONFIG_CONSTRUCTOR = configClass.getDeclaredConstructor()
+    CONFIG_CONSTRUCTOR.setAccessible(true)
+
     isConfigInstanceModifiable = true
   }
+
+  void cleanupSpec() {
+    rebuildConfig()
+  }
+
+  void setup() {
+    rebuildConfig()
+  }
+
+  void injectSysConfig(String name, String value) {
+    String prefixedName = name.startsWith("dd.") ? name : "dd." + name
+    System.setProperty(prefixedName, value)
+    rebuildConfig()
+  }
+
+  void injectEnvConfig(String name, String value) {
+    String prefixedName = name.startsWith("DD_") ? name : "DD_" + name
+    environmentVariables.set(prefixedName, value)
+    rebuildConfig()
+  }
+
+  /**
+   * Reset the global configuration. Please note that Runtime ID is preserved to the pre-existing value.
+   */
+  synchronized static void rebuildConfig() {
+    // Ensure the class was re-transformed properly in DDSpecification.makeConfigInstanceModifiable()
+    assert Modifier.isPublic(CONFIG_INSTANCE_FIELD.getModifiers())
+    assert Modifier.isStatic(CONFIG_INSTANCE_FIELD.getModifiers())
+    assert Modifier.isVolatile(CONFIG_INSTANCE_FIELD.getModifiers())
+    assert !Modifier.isFinal(CONFIG_INSTANCE_FIELD.getModifiers())
+
+    def newConfig = CONFIG_CONSTRUCTOR.newInstance()
+    CONFIG_INSTANCE_FIELD.set(null, newConfig)
+  }
 }
diff --git a/utils/test-utils/test-utils.gradle b/utils/test-utils/test-utils.gradle
@@ -6,4 +6,6 @@ dependencies {
 
   compile deps.bytebuddy
   compile deps.bytebuddyagent
+
+  compile group: 'com.github.stefanbirkner', name: 'system-rules', version: '1.19.0'
 }

Original file line number	Diff line number	Diff line change
`@@ -6,4 +6,6 @@ dependencies {`
`6`	`6`
`7`	`7`	`compile deps.bytebuddy`
`8`	`8`	`compile deps.bytebuddyagent`
	`9`	`+`
	`10`	`+ compile group: 'com.github.stefanbirkner', name: 'system-rules', version: '1.19.0'`
`9`	`11`	`}`