Fault Injection based on RandomizedRunner

FaultInjectionRunner.java

import com.carrotsearch.randomizedtesting.RandomizedRunner;
import org.jboss.byteman.contrib.bmunit.BMRule;
import org.jboss.byteman.contrib.bmunit.BMRules;
import org.jboss.byteman.contrib.bmunit.BMRunnerUtil;
import org.jboss.byteman.contrib.bmunit.BMScript;
import org.jboss.byteman.contrib.bmunit.BMScripts;
import org.jboss.byteman.contrib.bmunit.BMUnit;
import org.jboss.byteman.contrib.bmunit.BMUnitConfig;
import org.jboss.byteman.contrib.bmunit.BMUnitConfigState;
import org.junit.runners.model.InitializationError;
import org.junit.runners.model.Statement;

import java.lang.reflect.Method;

/**
 * Fault Injection Runner empowered by byteman.
 */
public class FaultInjectionRunner extends RandomizedRunner {

    private BMUnitConfig classConfigAnnotation;
    BMScript classSingleScriptAnnotation;
    BMScripts classMultiScriptAnnotation;
    BMRules classMultiRuleAnnotation;
    BMRule classSingleRuleAnnotation;
    Class<?> testKlazz;
    private final org.jboss.byteman.contrib.bmunit.BMRunnerUtil BMRunnerUtil = new BMRunnerUtil();

    /**
     * Creates a new runner for the given class.
     *
     * @param testClass
     */
    public FaultInjectionRunner(Class<?> testClass) throws InitializationError {
        super(testClass);
        testKlazz = testClass;
        classConfigAnnotation = testKlazz.getAnnotation(BMUnitConfig.class);
        classSingleScriptAnnotation = testKlazz.getAnnotation(BMScript.class);
        classMultiScriptAnnotation = testKlazz.getAnnotation(BMScripts.class);
        classMultiRuleAnnotation = testKlazz.getAnnotation(BMRules.class);
        classSingleRuleAnnotation = testKlazz.getAnnotation((BMRule.class));
        if (classMultiRuleAnnotation != null && classSingleRuleAnnotation != null) {
            throw new InitializationError("Use either BMRule or BMRules annotation but not both");
        }
        if (classMultiScriptAnnotation != null && classSingleScriptAnnotation != null) {
            throw new InitializationError("Use either BMScript or BMScripts annotation but not both");
        }
    }

    @Override
    protected Statement withClassRules(Statement s) {
        Statement statement = super.withClassRules(s);
        // n.b. we add the wrapper code in reverse order to the preferred order of loading
        // as it works by wrapping around and so execution is  in reverse order to wrapping
        // i.e. this ensures that the class script rules get loaded before any rules specified
        // using BMRule(s) annotations
        statement = addClassSingleRuleLoader(statement);
        statement = addClassMultiRuleLoader(statement);
        statement = addClassSingleScriptLoader(statement);
        statement = addClassMultiScriptLoader(statement);
        statement = addClassConfigLoader(statement);
        return statement;
    }

    protected Statement addClassConfigLoader(final Statement statement) {
        return new Statement() {
            public void evaluate() throws Throwable {
                BMUnitConfigState.pushConfigurationState(classConfigAnnotation, testKlazz);
                try {
                    statement.evaluate();
                } finally {
                    BMUnitConfigState.popConfigurationState(testKlazz);
                }
            }
        };
    }

    protected Statement addClassSingleScriptLoader(final Statement statement) {
        if (classSingleScriptAnnotation == null) {
            return statement;
        } else {
            final String name = BMRunnerUtil.computeBMScriptName(classSingleScriptAnnotation.value());
            final String loadDirectory = BMRunnerUtil.normaliseLoadDirectory(classSingleScriptAnnotation);
            return new Statement() {
                public void evaluate() throws Throwable {
                    BMUnit.loadScriptFile(testKlazz, name, loadDirectory);
                    try {
                        statement.evaluate();
                    } finally {
                        BMUnit.unloadScriptFile(testKlazz, name);
                    }
                }
            };
        }
    }

    protected Statement addClassMultiScriptLoader(final Statement statement) {
        if (classMultiScriptAnnotation == null) {
            return statement;
        } else {
            BMScript[] scriptAnnotations = classMultiScriptAnnotation.scripts();
            Statement result = statement;
            // note we iterate down here because we generate statements by wraparound
            // which means the the outer statement gets executed first
            for (int i = scriptAnnotations.length; i > 0; i--) {
                BMScript scriptAnnotation = scriptAnnotations[i - 1];
                final String name = BMRunnerUtil.computeBMScriptName(scriptAnnotation.value());
                final String loadDirectory = BMRunnerUtil.normaliseLoadDirectory(scriptAnnotation);
                final Statement nextStatement = result;
                result = new Statement() {
                    public void evaluate() throws Throwable {
                        BMUnit.loadScriptFile(testKlazz, name, loadDirectory);
                        try {
                            nextStatement.evaluate();
                        } finally {
                            BMUnit.unloadScriptFile(testKlazz, name);
                        }

                    }
                };
            }
            return result;
        }
    }

    protected Statement addClassMultiRuleLoader(final Statement statement) {
        if (classMultiRuleAnnotation == null) {
            return statement;
        } else {
            final String scriptText = BMRunnerUtil.constructScriptText(classMultiRuleAnnotation.rules());
            return new Statement() {
                public void evaluate() throws Throwable {
                    BMUnit.loadScriptText(testKlazz, null, scriptText);
                    try {
                        statement.evaluate();
                    } finally {
                        BMUnit.unloadScriptText(testKlazz, null);
                    }
                }
            };
        }
    }

    protected Statement addClassSingleRuleLoader(final Statement statement) {
        if (classSingleRuleAnnotation == null) {
            return statement;
        } else {
            final String scriptText = BMRunnerUtil.constructScriptText(new BMRule[]{classSingleRuleAnnotation});
            return new Statement() {
                public void evaluate() throws Throwable {
                    BMUnit.loadScriptText(testKlazz, null, scriptText);
                    try {
                        statement.evaluate();
                    } finally {
                        BMUnit.unloadScriptText(testKlazz, null);
                    }
                }
            };
        }
    }


    @Override
    protected Statement wrapMethodRules(Statement s, RandomizedRunner.TestCandidate c, Object instance) {
        Statement statement = super.wrapMethodRules(s, c, instance);
        // n.b. we add the wrapper code in reverse order to the preferred order of loading
        // as it works by wrapping around and so execution is in reverse order to wrapping
        // i.e. this ensures that the method script rules get loaded before any rules specified
        // using BMRule(s) annotations
        statement = addMethodSingleRuleLoader(statement, c.method);
        statement = addMethodMultiRuleLoader(statement, c.method);
        statement = addMethodSingleScriptLoader(statement, c.method);
        statement = addMethodMultiScriptLoader(statement, c.method);
        statement = addMethodConfigLoader(statement, c.method);
        return statement;
    }

    protected Statement addMethodConfigLoader(final Statement statement, Method method) {
        final BMUnitConfig annotation = method.getAnnotation(BMUnitConfig.class);
        return new Statement() {
            public void evaluate() throws Throwable {
                BMUnitConfigState.pushConfigurationState(annotation, method);
                try {
                    statement.evaluate();
                } finally {
                    BMUnitConfigState.popConfigurationState(method);
                }
            }
        };
    }

    /**
     * wrap the test method execution statement with the necessary
     * load and unload calls if it has a BMScript annotation
     * @param statement the statement to be evaluated
     * @param method the method being tested
     * @return the statement possibly wrapped with load and unload
     * calls
     */
    protected Statement addMethodSingleScriptLoader(final Statement statement, Method method) {
        BMScript annotation = method.getAnnotation(BMScript.class);
        if (annotation == null) {
            return statement;
        } else {
            // ensure we always have an actual name here instead of null because using
            // null will clash with the name used for looking up rules when the clas
            // has a BMRules annotation
            final String name = BMRunnerUtil.computeBMScriptName(annotation.value(), method);
            final String loadDirectory = BMRunnerUtil.normaliseLoadDirectory(annotation);
            return new Statement() {
                public void evaluate() throws Throwable {
                    BMUnit.loadScriptFile(testKlazz, name, loadDirectory);
                    try {
                        statement.evaluate();
                    } finally {
                        BMUnit.unloadScriptFile(testKlazz, name);
                    }
                }
            };
        }
    }

    /**
     * wrap the test method execution statement with the necessary
     * load and unload calls if it has a BMScripts annotation
     * @param statement the statement to be evaluated
     * @param method the method being tested
     * @return the statement possibly wrapped with load and unload
     * calls
     */
    protected Statement addMethodMultiScriptLoader(final Statement statement, Method method) {
        BMScripts scriptsAnnotation = method.getAnnotation(BMScripts.class);
        if (scriptsAnnotation == null) {
            return statement;
        } else {
            BMScript[] scriptAnnotations = scriptsAnnotation.scripts();
            Statement result = statement;
            // note we iterate down here because we generate statements by wraparound
            // which means the the outer statement gets executed first
            for (int i = scriptAnnotations.length; i > 0; i--) {
                BMScript scriptAnnotation = scriptAnnotations[i - 1];
                final Statement nextStatement = result;
                // ensure we always have an actual name here instead of null because using
                // null will clash with the name used for looking up rules when the clas
                // has a BMRules annotation
                final String name = BMRunnerUtil.computeBMScriptName(scriptAnnotation.value(), method);
                final String loadDirectory = BMRunnerUtil.normaliseLoadDirectory(scriptAnnotation);
                result = new Statement() {
                    public void evaluate() throws Throwable {
                        BMUnit.loadScriptFile(testKlazz, name, loadDirectory);
                        try {
                            nextStatement.evaluate();
                        } finally {
                            BMUnit.unloadScriptFile(testKlazz, name);
                        }
                    }
                };
            }
            return result;
        }
    }

    /**
     * wrap the test method execution statement with the necessary
     * load and unload calls if it has a BMRules annotation
     * @param statement the statement to be evaluated
     * @param method the method being tested
     * @return the statement possibly wrapped with load and unload
     * calls
     */
    protected Statement addMethodMultiRuleLoader(final Statement statement, Method method) {
        BMRules annotation = method.getAnnotation(BMRules.class);
        if (annotation == null) {
            return statement;
        } else {
            final String name = method.getName();
            final String script = BMRunnerUtil.constructScriptText(annotation.rules());
            return new Statement() {
                public void evaluate() throws Throwable {
                    BMUnit.loadScriptText(testKlazz, name, script);
                    try {
                        statement.evaluate();
                    } finally {
                        BMUnit.unloadScriptText(testKlazz, name);
                    }
                }
            };
        }
    }

    /**
     * wrap the test method execution statement with the necessary
     * load and unload calls if it has a BMRule annotation
     * @param statement the statement to be evaluated
     * @param method the method being tested
     * @return the statement possibly wrapped with load and unload
     * calls
     */
    protected Statement addMethodSingleRuleLoader(final Statement statement, Method method) {
        BMRule annotation = method.getAnnotation(BMRule.class);
        if (annotation == null) {
            return statement;
        } else {
            final String name = method.getName();
            final String script = BMRunnerUtil.constructScriptText(new BMRule[]{annotation});
            return new Statement() {
                public void evaluate() throws Throwable {
                    BMUnit.loadScriptText(testKlazz, name, script);
                    try {
                        statement.evaluate();
                    } finally {
                        BMUnit.unloadScriptText(testKlazz, name);
                    }
                }
            };
        }
    }

}

TestLucene70DocValuesFormat.java

// Lucene example

@RunWith(FaultInjectionRunner.class)
public class TestLucene70DocValuesFormat extends BaseCompressingDocValuesFormatTestCase {
  
    @BMRule(name = "index writer failure",
            targetClass = "org.apache.lucene.index.IndexWriter",
            targetMethod = "addDocuments",
            action = "throw new IOException(\"throwing io exception\")")
    public void testTermsEnumLongSharedPrefixes() throws Exception {
    int numIterations = atLeast(1);
    for (int i = 0; i < numIterations; i++) {
      doTestTermsEnumRandom(TestUtil.nextInt(random(), 1025, 5121), () -> {
        char[] chars = new char[random().nextInt(500)];
        Arrays.fill(chars, 'a');
        if (chars.length > 0) {
          chars[random().nextInt(chars.length)] = 'b';
        }
        return new String(chars);
      });
    }
  }
  
  ...
  
}