build.please.cover.runner.InstrumentingClassLoader Maven / Gradle / Ivy

Go to download

Show more of this group Show more artifacts with this name
Show all versions of junit-runner Show documentation

JUnit test runner for Please to run Java tests with.

The newest version!

package build.please.cover.runner;

import build.please.vendored.org.jacoco.core.instr.Instrumenter;

import java.io.IOException;
import java.io.InputStream;
import java.util.HashMap;
import java.util.Map;
import java.util.Set;

/**
 * Loads and instruments classes for coverage.
 */
public class InstrumentingClassLoader extends ClassLoader {
  private final Instrumenter instrumenter;
  private final Map> instrumentedClasses = new HashMap<>();

  InstrumentingClassLoader(Instrumenter instrumenter) {
    this.instrumenter = instrumenter;
  }

  void addInstrumentedClasses(Set classes) {
    for (String cls : classes) {
        instrumentedClasses.put(cls, null);
    }
  }

  Iterable> getInstrumentedClasses() {
    return instrumentedClasses.values();
  }

  @Override
  protected Class loadClass(String name, boolean resolve) throws ClassNotFoundException {
    Class c = findLoadedClass(name);
    if (c != null) {
      return c;
    }

    try {
      Class cls = instrumentedClasses.get(name);
      if (cls != null) {
        return cls;
      } else if (instrumentedClasses.containsKey(name)) {
        byte[] instrumented = instrumenter.instrument(getTargetClass(InstrumentingClassLoader.class, name), name);
        cls = defineClass(name, instrumented, 0, instrumented.length, this.getClass().getProtectionDomain());
        instrumentedClasses.put(name, cls);
        return cls;
      }
      return super.loadClass(name, resolve);
    } catch (IOException ex) {
      throw new RuntimeException(ex);
    }
  }

  static InputStream getTargetClass(Class cls, String name) {
    final String resource = '/' + name.replace('.', '/') + ".class";
    return cls.getResourceAsStream(resource);
  }
}