Reset class static variable during unit test

Question

I am trying to write a unit test for a legacy code. The class which I'm testing has several static variables. My test case class has a few @Test methods. Hence all of them share the same state.

Is there way to reset all static variables between tests?

One solution I came up is to explicitly reset each field, e.g.:

field(MyUnit.class, "staticString").set(null, null);
((Map) field(MyUnit.class, "staticFinalHashMap").get(null)).clear();

As you see, each variable needs custom re-initialization. The approach is not easy to scale, there are a lot such classes in the legacy code base. Is there any way to reset everything at once? Maybe by reloading the class each time?

As a possible good solution I think is to use something like powermock and create a separate classloader for each test. But I don't see easy way to do it.

Solution 1

Ok, I think I figured it out. It is very simple.

It is possible to move @PrepareForTest powermock's annotation to the method level. In this case powermock creates classloader per method. So it does that I need.

Solution 2

Let's say I'm testing some code involving this class:

import java.math.BigInteger;
import java.util.HashSet;

public class MyClass {
  static int someStaticField = 5;
  static BigInteger anotherStaticField = BigInteger.ONE;
  static HashSet<Integer> mutableStaticField = new HashSet<Integer>();
}

You can reset all of the static fields programmatically using Java's reflection capabilities. You will need to store all of the initial values before you begin the test, and then you'll need to reset those values before each test is run. JUnit has @BeforeClass and @Before annotations that work nicely for this. Here's a simple example:

import static org.junit.Assert.*;

import java.lang.reflect.Field;
import java.math.BigInteger;
import java.util.Map;
import java.util.HashMap;

import org.junit.AfterClass;
import org.junit.Before;
import org.junit.BeforeClass;
import org.junit.Test;

public class MyTest extends Object {

  static Class<?> staticClass = MyClass.class;
  static Map<Field,Object> defaultFieldVals = new HashMap<Field,Object>();

  static Object tryClone(Object v) throws Exception {
    if (v instanceof Cloneable) {
      return v.getClass().getMethod("clone").invoke(v);
    }
    return v;
  }

  @BeforeClass
  public static void setUpBeforeClass() throws Exception {
    Field[] allFields = staticClass.getDeclaredFields();
    try {
      for (Field field : allFields) {
          if (java.lang.reflect.Modifier.isStatic(field.getModifiers())) {
              Object value = tryClone(field.get(null));
              defaultFieldVals.put(field, value);
          }
      }
    }
    catch (IllegalAccessException e) {
      System.err.println(e);
      System.exit(1);
    }
  }

  @AfterClass
  public static void tearDownAfterClass() {
    defaultFieldVals = null;
  }

  @Before
  public void setUp() throws Exception {
    // Reset all static fields
    for (Map.Entry<Field, Object> entry : defaultFieldVals.entrySet()) {
      Field field = entry.getKey();
      Object value = entry.getValue();
      Class<?> type = field.getType();
      // Primitive types
      if (type == Integer.TYPE) {
        field.setInt(null, (Integer) value);
      }
      // ... all other primitive types need to be handled similarly
      // All object types
      else {
        field.set(null, tryClone(value));
      }
    }
  }

  private void testBody() {
    assertTrue(MyClass.someStaticField == 5);
    assertTrue(MyClass.anotherStaticField == BigInteger.ONE);
    assertTrue(MyClass.mutableStaticField.isEmpty());
    MyClass.someStaticField++;
    MyClass.anotherStaticField = BigInteger.TEN;
    MyClass.mutableStaticField.add(1);
    assertTrue(MyClass.someStaticField == 6);
    assertTrue(MyClass.anotherStaticField.equals(BigInteger.TEN));
    assertTrue(MyClass.mutableStaticField.contains(1));
  }

  @Test
  public void test1() {
    testBody();
  }

  @Test
  public void test2() {
    testBody();
  }

}

As I noted in the comments in setUp(), you'll need to handle the rest of the primitive types with similar code for that to handle ints. All of the wrapper classes have a TYPE field (e.g. Double.TYPE and Character.TYPE) which you can check just like Integer.TYPE. If the field's type isn't one of the primitive types (including primitive arrays) then it's an Object and can be handled as a generic Object.

The code might need to be tweaked to handle final, private, and protected fields, but you should be able to figure how to do that from the documentation.

Good luck with your legacy code!

Edit:

I forgot to mention, if the initial value stored in one of the static fields is mutated then simply caching it and restoring it won't do the trick since it will just re-assign the mutated object. I'm also assuming that you'll be able to expand on this code to work with an array of static classes rather than a single class.

Edit:

I've added a check for Cloneable objects to handle cases like the HashMap in your example. Obviously it's not perfect, but hopefully this will cover most of the cases you'll run in to. Hopefully there are few enough edge cases that it won't be too big of a pain to reset them by hand (i.e. add the reset code to the setUp() method).

Comments

[kan]

It will work until you have a mutable value. E.g. try to do the trick with a HashMap instead of BigInteger.

[DaoWen]

I've updated the code to call clone() on the default value if it's Cloneable. That should help... If that doesn't do it for you then you can either handle the edge cases manually or go with the suggestion of using a separate class loader for each of the objects with static fields.

[DaoWen]

This seems way nicer than my terrible reflection hack. Good to know this exists!

[Pankaj Vatsa]

Thanks for this answer!

[Jeremy Kahan]

Having a hard time making this work in 2021. Does it still? (junit 4, jupiter api 5.7) In the test file, say I have Test1 and Test2. You would say @PrepareForTest(Main.class) before/after the @test annotation for each? Or would it go with BeforeEach?