Search code examples
javajunitjavassist

Javassist API with Junit


I am attempting to change some third party class definitions, before each test, to simulate different results. I have to use something like javassist because extending the classes, sometimes, is just not possible due to the access modifiers. Here is an example of what I am attempting to do with javassist and junit combined:

public class SimulatedSession extends SomeThirdParty {

    private boolean isJoe = false;
    public SimulatedSession(final boolean isJoe) {
        this.isJoe = isJoe;
    }

    @Override
    public void performThis() {
        final ClassPool classPool = ClassPool.getDefault();
        final CtClass internalClass = classPool.get("some.package.Class");
        final CtMethod callMethod = internalClass.getDeclaredMethod("doThis");
        if (isJoe) {
            callMethod.setBody("{System.out.println(\"Joe\");}");
        } else {
            callMethod.setBody("{System.out.println(\"mik\");}");
        }
        internalClass.toClass();
    }
}

@Test
public void firstTest() {
     SimulatedSession toUse = new SimulatedSession(false);
     // do something with this object and this flow
}

@Test
public void nextTest() {
     SimulatedSession toUse = new SimulatedSession(true);
     // do something with this object and this flow
}

if I run each test individually, I can run the code just fine. When I run them using the unit suite, one test after the other, I get a "frozen class issue". To get around this, I am looking at this post, however, I must admit I am unsure as to how one can use a different class pool to solve the issue.


Solution

  • Your current code will try to load twice the same class into the same ClassLoader which is forbidden, you can only load once a class for a given ClassLoader.

    To make your unit tests pass, I had to:

    1. Create my own temporary ClassLoader that will be able to load some.package.Class (that I replaced by javassist.MyClass for testing purpose) and that will be implemented in such way that it will first try to load the class from it before the parent's CL.
    2. Set my own ClassLoader as context ClassLoader.
    3. Change the code of SimulatedSession#performThis() to be able to get the class instance created by this method and to call internalClass.defrost() to prevent the "frozen class issue".
    4. Invoke by reflection the method doThis() to make sure that I have different output by using the class instance returned by SimulatedSession#performThis() to make sure that the class used has been loaded with my ClassLoader.

    So assuming that my class javassist.MyClass is:

    package javassist;
    
    public class MyClass {
        public void doThis() {
    
        }
    }
    

    The method SimulatedSession#performThis() with the modifications:

    public Class<?> performThis() throws Exception {
        final ClassPool classPool = ClassPool.getDefault();
        final CtClass internalClass = classPool.get("javassist.MyClass");
        // Prevent the "frozen class issue"
        internalClass.defrost();
        ...
        return internalClass.toClass();
    }
    

    The unit tests:

    // The custom CL
    private URLClassLoader cl;
    // The previous context CL
    private ClassLoader old;
    
    @Before
    public void init() throws Exception {
        // Provide the URL corresponding to the folder that contains the class
        // `javassist.MyClass`
        this.cl = new URLClassLoader(new URL[]{new File("target/classes").toURI().toURL()}){
            protected Class<?> loadClass(String name, boolean resolve)
                throws ClassNotFoundException {
                try {
                    // Try to find the class for this CL
                    return findClass(name);
                } catch( ClassNotFoundException e ) {
                    // Could not find the class so load it from the parent
                    return super.loadClass(name, resolve);
                }
            }
        };
        // Get the current context CL and store it into old
        this.old = Thread.currentThread().getContextClassLoader();
        // Set the custom CL as new context CL
        Thread.currentThread().setContextClassLoader(cl);
    }
    
    @After
    public void restore() throws Exception {
        // Restore the context CL
        Thread.currentThread().setContextClassLoader(old);
        // Close the custom CL
        cl.close();
    }
    
    
    @Test
    public void firstTest() throws Exception {
        SimulatedSession toUse = new SimulatedSession(false);
        Class<?> c = toUse.performThis();
        // Invoke doThis() by reflection
        Object o2 = c.newInstance();
        c.getMethod("doThis").invoke(o2);
    }
    
    @Test
    public void nextTest() throws Exception {
        SimulatedSession toUse = new SimulatedSession(true);
        Class<?> c = toUse.performThis();
        // Invoke doThis() by reflection
        Object o2 = c.newInstance();
        c.getMethod("doThis").invoke(o2);
    }
    

    Output:

    mik
    Joe