06 June 2011

@PublicForTest letting tests access private members

Yesterday, I was looking at the JavaDocs for Google Guava (what else would I do with my weekend?).  Specifically I was looking at  the @VisibleForTesting annotation, and thinking that it was a shame that class members can't be kept private for production code, and still be usable by tests.

Then it occurred to me that Project Lombok has proved that the code you type doesn't have to be exactly the code that the compiler sees.  Now, I'm not entirely comfortable with that much magic going on in production code, but in test code I have more control over the environment, and users won't be inconvenienced if it fails.

The Concept
So, my idea is to use java.lang.instrument, and BCEL to make a java agent, which would find methods and fields annotated with a new annotation: @PublicForTest, and change their visibility to make them public.  Then, this java agent could be used when running unit tests, and tests could be written that directly accessed non-public fields and methods.

Problems
The current version (5.2) of BCEL doesn't give you any help when it comes to annotations.  The next version (5.3) looks like it has support for querying annotations, but I'd rather not touch it while it's a snapshot version.  Fortunately, I came across an article from a few years ago by Don Schwarz, with examples of using BCEL with annotations.

Progress So Far
A few hours later, I have a java agent that will make private turn into public.  So, now if I could compile a test that accessed a private member, it would run.  But, of course I can't compile code that accesses a private field/method.  So, the next task is to trick the compiler into thinking the private fields are public.

Tricking the Compiler
The java agent I've written so far won't do it, because the compiler isn't loading the classes, it's reading them as files.  I don't believe an annotation processor would do it, because I don't want this to affect the class file created that uses the @PrivateForTest annotation, I just want to affect how the compiler sees the already compiled class when it's compiling the test classes.   I've had a look inside javac, and it looks like all file access goes through implementations of an interface called JavaFileManager.  I had hoped to discover a hidden option to the compiler to allow me to push in my own implementation of this, but no luck.  My current idea is to add to the java agent to do something similar to an AOP around advice.  If the java agent transforms any subclass of JavaFileManager to return my own implementation of JavaFileObject or FileObject, then I should be able to transform the class information that the compiler reads, again turning private turn into public.

This meddling with the compiler internals is very ugly, but I'm hoping a little ugly code in one place may be worth it to make things nicer elsewhere.  I'd be interested if anyone has better ideas on how to achieve the same ends.

The Code
This was my first attempt with using either BCEL, or java.lang.instrument, so I won't post everything until I've got it tidied up, and have the compile-time side of things figured out.  But, here's the basics:

The annotation PublicForTest.java
public @interface PublicForTest {
}

META-INF/MANIFEST.MF
I actually configured maven to add this line to the manifest.  When you run java using -javaagent: it looks in the manifest of the jar you give it for a line like the following to tell it where the agent's main class is.
Premain-Class: org.upbiit.publicfortest.PreMain

The main class PreMain.java
All my PreMain class does is register a transformer to use on any classes that get loaded.
public class PreMain {
  public static void premain(String agentArgs, Instrumentation inst) {
    inst.addTransformer( new PublicForTestTransformer(agentArgs));
  }
}

The transformer PublicForTestTransformer.java
Once a ClassFileTransformer has been registered with Instrumentation, this class will be called for every class that gets loaded, it has the opportunity to change the bytes that make up the class file. Here, I've used BCEL to read the class, modify the visibility of the fields and methods that I'm interested in, and turn it back into bytes.
class PublicForTestTransformer implements ClassFileTransformer {
  static {
    AnnotationReader ar = new AnnotationReader();
    Attribute.addAttributeReader( "RuntimeVisibleAnnotations", ar );
    Attribute.addAttributeReader( "RuntimeInvisibleAnnotations", ar );
  }
  public PublicForTestTransformer( String agentArgs ) {
    //do nothing
  }
  public byte[] transform(ClassLoader loader,
                 String className,
                 Class classBeingRedefined,
                 ProtectionDomain protectionDomain,
                 byte[] classfileBuffer)
                 throws IllegalClassFormatException {
    try {
    ClassParser parser = new ClassParser(
        new ByteArrayInputStream( classfileBuffer), className+".java"
    );
    JavaClass javaClass = parser.parse();
    ClassGen classGen = new ClassGen( javaClass );

    for(Method m : classGen.getMethods()) {
      if( hasPFTAnn(m)) {
        m.isPrivate(false);
        m.isProtected( false );
        m.isPublic(true);
      }
    }
    for(Field f : classGen.getFields()) {
      if( hasPFTAnn(f)) {
        f.isPrivate(false);
        f.isProtected( false );
        f.isPublic(true);
      }
    }
    return classGen.getJavaClass().getBytes();
    } catch( IOException e ) {
      throw new RuntimeException(e);
    }
  }
  //...
}

Test Code
A class with private members
public class ClassWithPrivates {
  @PublicForTest
  private int privateField;

  @PublicForTest
  private int privateMethod() {
    return privateField;
  }
  @PublicForTest
  private void privateMethod(int i) {
    privateField = i;
  }
}

A class that checks the members of ClassWithPrivates
public class Test {
    public static void main(String[] args) {
      ClassWithPrivates n = new ClassWithPrivates();
      //Commented out until the compiler can be tricked into accepting access 
      //to private members
      //n.privateField = 3;
      //assert n.privateField == 3;

      try {
        System.out.println("Fields");
        System.out.println(ClassWithPrivates.class.getField("privateField"));

        System.out.println("Methods");
        for( Method m : ClassWithPrivates.class.getDeclaredMethods()) {
          System.out.println(m);
        }
      } catch( Exception e) {
        throw new RuntimeException(e);
      }
    }

}

Output from Test
Note that it says they are public, even though they were declared private. This is because I ran Test with my publicfortest.jar java agent.
Fields
public int org.upbiit.publicfortest.NewClass.privateField
Methods
public int org.upbiit.publicfortest.NewClass.privateMethod()
public void org.upbiit.publicfortest.NewClass.privateMethod(int)

Caveats
There may be a case to be made that the changes created by this tool would reduce the validity of tests, as the class under test would be different than it is in production. However, the changes are limited, and with mocks/stubs/etc. we have differences between test and production already.
As long as this tool isn't applied to the compiler when it is compiling non-test code, there should be no risk to code that isn't using reflection.
In the case of production code using reflection, there would be a risk that during production with this tool turned off it may fail to access fields/methods that it was able to access during test.

Possible Future Work
IDEs
  • Once the compiler and the runtime can both be convinced to treat members with @PublicForTest as public, this code will be usable for making tests more readable, without compromising the production code when it's not under test. But, IDEs and editors may still complain about the tests' attempts to access private members.

Other compilers
  • The next lot of code to trick the compiler into accepting access to private members, will be a bit of a hack. I expect it would only work with the Oracle JDK, and probably also with OpenJDK.

@PubicFor(?)
  • What if there are some members that you want to make accessible for unit tests and larger tests, and some that should only be made accessible for unit tests?
    If the @PublicForTest annotation was extended to allow a keyword that would also have to be supplied to the java agent, then this could be controlled.
    e.g.
    @PublicForTest("unit")
    or
    @PublicForTest({"unit","integration"})
    -javaagent:publicfortest.jar=scope=unit

Restrict packages
  • While you may want to be able to access private members of your own code under test, if you used a library that also used @PublicForTest for testing, its private members probably shouldn't be accessible. So, it may make sense to give the java agent an argument telling it to only pay attention to classes whose full name match a regex, or are in a given package.
    e.g.
    -javaagent:publicfortest.jar=in=org.upbiit.mypackage

Adding support for other compilers, and for IDEs is more work than I personally would find worthwhile.  I may consider it if I got a lot of interest.

Summary
BCEL proved to be pretty straight-forward to use, but I wasn't doing anything complicated with it.  When I start using it to change the innards of the compiler, then I'll have to actually create instructions with it.  This is where I expect it will get more complicated (but only in the way that writing assembly code is more complicated than writing source code).

There may be some risk that, because of this tool, code under test could work where code under production would fail. But, this risk is often there with unit tests, as the class under test is using mocks/stubs/etc. instead of the real thing.

When I'm done, I'll post all of the code on my site under an Apache licence.

It might be ugly, but I hope that when it's done, this code might make the boundary between unit tests and production code a little prettier.