Getting rid of dependcies (part 2) - Stubbing singletons

In the previous post about removing dependencies I used approach which boils down to adding a layer of indirection in the code. (By dependencies I here mean "evil" dependencies on a classes which need some external resource(s) to be created or used.) Unfortunately it may not always be possible to add lots of changes to the source code. For me the main reason was that other team members weren't looking forward to it. But there may be other political reasons to avoid changes throughout you code base. For example you may be asked to add a relatively small feature into the project you don't own and there may be no definite plans to make you the owner. Though the are no technical reasons, I can think of, to avoid changing code if are plagued with dependencies.

Messy in some places the code-base I work with has some good design decisions. One of them is that most of the access to external resources is done through a certain group of classes. Most of these classes are singletons. And if a class itself is not a singleton then the factory that creates it is a singleton or at least it's accessed through a facade which is a singleton. In all the cases, I came across, somewhere on the way of accessing external resource there is a singleton involved. But aren't singletons evil? As it turned out from dependency removing point of view singletons may be not really bad. The good part of is that they have one instance, and if you somehow managed to replace this "only" instance with your implementation, then you change behaviour of user code that use singleton without making any changes in it. To a degree it looks like turning getInstance() method from getter into a factory method. Singleton replacing can be though of as two tasks:

• subclassing singleton;
• replacing singleton's "only" instance.

Below are possible variations of doing these tasks depending on singleton implementation. I didn't actually use all of the described on the real system. The sole reason I explored possibilities was to try and stretch limits of what seems to be doable.

Stubbing interface-based singleton with lazy initialization

By interface-based singleton I mean a singleton like this one (if someone cares DataMunger class name is inspired by this kata :) ):

public class DefaultDataMunger implements DataMunger {
    private static DataMunger singleInstance;

    public static DataMunger getInstance() {
        if (singleInstance == null) {
            singleInstance = new DefaultDataMunger();
        }
        return singleInstance;
    }

    private DefaultDataMunger() {}

    public String munge(final Data data) {
        return "I access external resource...";
    }
}

public interface DataMunger {
    String munge(Data data);
}

Subclassing singleton in this case is easy since singleton implements an interface and getInstance() method returns it. We can just create stub implementation of this interface. Replacing singleton instance is not much harder. This can be easily done with reflection. On the whole a unit test for stubbing might look like this:

public class SingletonRemovingTest ...
    private static final Data DONT_CARE_WHAT = null;

    @Test public void shouldSubstituteInterfaceBasedSingleton() throws Exception {
        // setup
        DataMunger dataMunger = new DataMunger() {
            public String munge(final Data data) {
                return "I'm a simple implementation";
            }
        };
        setPrivateField(DefaultDataMunger.class, "singleInstance", dataMunger);

        // exercise / verify
        assertEquals("I'm a simple implementation", DefaultDataMunger.getInstance().munge(DONT_CARE_WHAT));
    }

    private static void setPrivateField(final Class aClass, final String fieldName, final Object object)
            throws NoSuchFieldException, IllegalAccessException {
        final Field field = aClass.getDeclaredField(fieldName);
        field.setAccessible(true);
        field.set(null, object);
    }

This is perhaps the simplest kind of singleton to stub.

Stubbing class-based singletons with lazy initialization

This kind of singleton looks like this:

public class DataProcessor {
    private static DataProcessor singleInstance;

    public static DataProcessor getInstance() {
        if (singleInstance == null) {
            singleInstance = new DataProcessor();
        }
        return singleInstance;
    }

    private DataProcessor() {
    }

    public String retrieveUpdate() {
        return "update from far away";
    }

    public Data applyUpdate(final Data data, final String update) {
        return null;
    }
}

In this case singleton's getInstance() return type is the singleton itself. Though it's still quite easy to replace singleInstance, there is a problem with subclassing because the only constructor of singleton is private. The simple workaround is to use some mocking library which is capable of mocking classes (here mockito is used, but jmock or easymock will do as well):

    @Test public void shouldSubstituteClassBasedSingleton() throws Exception {
        // setup
        DataProcessor dummyProcessor = mock(DataProcessor.class);
        stub(dummyProcessor.retrieveUpdate()).toReturn("dummy update");
        setPrivateField(DataProcessor.class, "singleInstance", dummyProcessor);

        // exercise / verify
        assertEquals("dummy update", DataProcessor.getInstance().retrieveUpdate());
    }

Here stub() method sets new behaviour for the retrieveUpdate() method. To some extent it can be thought of as overriding and the whole thing like crippled subclassing. (It's possible to make it even more like subclassing by using toAnswer() method and passing in it implementation of Answer interface with whatever code you need.)

Stubbing final class-based singletons with lazy initialization

This type of singleton is almost the same as the previous one with the difference that it has final modifier in declaration. And there is a problem with it since it cannot be subclassed. Among mocking java frameworks I worked with jmock, easymock and mockito (well, it uses easymock inside but anyway) and in order to stub a class all of them create a subclass. It basically means that you can't mock a final class because you need to subclass it (at least that's how I understand it). So, for instance, if you try to run this test with mockito

    @Test public void shouldSubstituteFinalClassBasedSingleton() throws Exception {
        // setup
        FinalDataProcessor dummyProcessor = mock(FinalDataProcessor.class);
        stub(dummyProcessor.retrieveUpdate()).toReturn("dummy update");
        setPrivateField(FinalDataProcessor.class, "singleInstance", dummyProcessor);

        // exercise / verify
        assertEquals("dummy update", FinalDataProcessor.getInstance().retrieveUpdate());
    }

an error like this is inevitable:

    org.mockito.exceptions.base.MockitoException:
    Mockito cannot mock final classes like:
    class home.lib.FinalDataProcessor
        at home.v2.SingletonRemovingTest.shouldSubstituteFinalClassBasedSingleton(SingletonRemovingTest.java:43)
        ...

I found two possible ways to get around this problem. The first one is simplistic and works only if you have sources. It is to remove finals in the source code and recompile it. The disadvantage of it is that because you presumably don't own the code, you cannot commit changes you've made. Another approach is to remove finals at run-time. I believe there must be some project that removes finals, but I couldn't find anything alike so I rolled up my own not very ingenious solution. It is to setup a custom ClassLoader and fix class structure with byte-code manipulating tool.

Skipping horrid classloading details I simply replaced system classloader (the one which loads classes from the classpath):

public class ClassLoaderForTesting extends ClassLoader {
    private static List<String> CLASSES_TO_FIX = Arrays.asList(
            "home.lib.FinalDataProcessor",
            "home.lib.FinalStaticDataProcessor"
    );

    /**
     * @param classLoader parent classloader (it will be "original" system classloader)
     */
    public ClassLoaderForTesting(ClassLoader classLoader) {
        super(classLoader);
    }

    @Override protected synchronized Class<?> loadClass(final String name, final boolean resolve)
            throws ClassNotFoundException {
        if (isSystemClass(name)) {
            return super.loadClass(name, resolve);
        }
        return loadClassManually(name);
    }

    private Class<?> loadClassManually(final String className) throws ClassNotFoundException {
            byte[] classAsBytes = getClassAsBytes(className);
        if (CLASSES_TO_FIX.contains(className)) {
            classAsBytes = fixClassStructure(classAsBytes);
        }
        return defineClass(null, classAsBytes, 0, classAsBytes.length);
    }

This classloader loads all the classes found on the classpath fixing some of them along the way. Actual replacing of system classloader is done by specifying JVM parameter:

     -Djava.system.class.loader=home.v2.ClassLoaderForTesting 

As a byte-code fixing tool I used asm because it has great documentation and I used it before. The whole final-removing class looks like this:

public class FinalRemover {
    public byte[] process(final byte[] input) {
        ClassWriter classWriter = new ClassWriter(0);
        ClassAdapter classAdapter = new ClassAdapter(classWriter) {
            @Override public void visit(final int version, final int access, final String name, final String signature,
                                        final String superName, final String[] interfaces) {
                // remove final from class header
                super.visit(version, removeFinalFrom(access), name, signature, superName, interfaces);
            }

            @Override public FieldVisitor visitField(final int access, final String name, final String desc,
                                                     final String signature, final Object value) {
                // remove final from fields headers
                return super.visitField(removeFinalFrom(access), name, desc, signature, value);
            }
        };
        ClassReader classReader = new ClassReader(input);
        classReader.accept(classAdapter, 0);
        return classWriter.toByteArray();
    }

    private static int removeFinalFrom(final int access) {
        return access & (~Opcodes.ACC_FINAL);
    }
}

The class may seem a bit daunting, but it's not so. It transfers bytes from reader to writer changing some of them along the way with help of ClassAdapter. ( bytes --> ClassReader --> ClassAdapter --> ClassWriter --> bytes ) All the above code put together makes it possible to stub final singletons, so the above unit test passes even though at the source code level the class it stubs is final.

Stubbing final class-based singletons with eager initialization

This singleton is like previous ones but its only instance is created right away on the class initialization. Replacing the "only" instance with reflection won't help, since we are in trouble as soon as the class is loaded:

public final class FinalStaticDataProcessor {
    private final static FinalStaticDataProcessor singleInstance = new FinalStaticDataProcessor();

    public static FinalStaticDataProcessor getInstance() {
        return singleInstance;
    }

    private FinalStaticDataProcessor() {
        throw new IllegalStateException("long initialization process happened");
    }

    public String retrieveUpdate() {
        return "update from far away";
    }
}

    @Test public void shouldSubstituteFinalStaticClassBasedSingleton() throws Exception {
        // setup
        FinalStaticDataProcessor dummyProcessor = mock(FinalStaticDataProcessor.class);
        stub(dummyProcessor.retrieveUpdate()).toReturn("dummy update");
        setPrivateField(FinalStaticDataProcessor.class, "singleInstance", dummyProcessor);

        // exercise / verify
        assertEquals("dummy update", FinalStaticDataProcessor.getInstance().retrieveUpdate());
    }

If you run the above unit test with definalizing classloader, you will apparently get IllegalStateException. No surprises with the solution. The trick used for replacing finals is perfectly applicable here as well. So the whole fixClassStructure() method in ClassLoaderForTesting will be:

    private static byte[] fixClassStructure(byte[] classAsBytes) {
        classAsBytes = new FinalRemover().process(classAsBytes);
        classAsBytes = new StaticInitRemover().process(classAsBytes);
        return classAsBytes;
    }

I won't put StaticInitRemover class here since it's a lot of code but the idea is the same as with removing finals. StaticInitRemover watches byte-code in static initializers and hides from ClassWriter anything that resembles this:

    NEW home/lib/FinalStaticDataProcessor
    DUP
    INVOKESPECIAL home/lib/FinalStaticDataProcessor.<init> ()V
    PUTSTATIC home/lib/FinalStaticDataProcessor.singleInstance : Lhome/lib/FinalStaticDataProcessor;

So what?

The main lesson I learned is that it's not unimaginable to stub/mock final class or change value of a final field if you really want to. And most importantly it's not very hard to do so. Another thing is that there are actually no limits once you are on the road of byte-code changing. I imagine it's not undoable in a reasonable time to replace bad static invocations with invocations of stubbed methods or get rid of bad constructors' invocations. On other hand there is a point after which it doesn't make any sense to go further, since class changing trick is nothing but a workaround and the real problem is in the source code.

Update: a bit later I came across a post with more standard way to change classes upon loading. It is to use javaagent (see java.lang.instrument).