Mock janet commands

BaseTest.java

package com.ghm.util;

import com.ghm.TestGHMApplication;
import com.ghm.di.AppComponent;
import com.ghm.di.AppModule;
import com.ghm.di.MockCommandService;
import com.ghm.di.TestApiModule;
import com.ghm.di.TestJanetModule;
import com.ghm.service.api.CreatePSTNCallAction;

import org.robolectric.RuntimeEnvironment;

import java.util.ArrayList;
import java.util.List;

import io.techery.janet.ActionService;
import io.techery.janet.Command;
import rx.Scheduler;
import rx.plugins.RxJavaPlugins;
import rx.plugins.RxJavaSchedulersHook;
import rx.schedulers.Schedulers;

import static io.techery.presenta.mortar.DaggerService.createComponent;

public class BaseTest {

    protected TestGHMApplication application;
    private AppComponent component;
    private MockCommandService mockCommandService;

    public void setUp() {
        application = (TestGHMApplication) RuntimeEnvironment.application;
        component = createComponent(AppComponent.class, getModules().toArray());
        application.setComponent(component);

        for (ActionService service: component.provideServices()) {
            if (service instanceof MockCommandService) {
                mockCommandService = (MockCommandService) service;
            }
        }
    }

    protected void mockFail(Class<? extends Command> commandClass) {
        mockCommandService.mockFail(commandClass);
    }

    protected void mockSuccess(Class<? extends Command> commandClass) {
        mockCommandService.mockSuccess(commandClass);
    }

    protected void mockSuccess(Class<CreatePSTNCallAction> commandClass, Object result) {
        mockCommandService.mockSuccess(commandClass, result);
    }

    public List<Object> getModules() {
        ArrayList<Object> modules = new ArrayList<>();
        modules.add(new AppModule(application));
        modules.add(new TestApiModule());
        modules.add(new TestJanetModule());
        return modules;
    }

    public AppComponent getAppComponent() {
        return component;
    }

    static {
        RxJavaPlugins.getInstance().registerSchedulersHook(new RxJavaSchedulersHook() {
            @Override
            public Scheduler getIOScheduler() {
                return Schedulers.immediate();
            }
        });
    }
}

CommandUtils.java

package io.techery.janet;

import io.techery.janet.Command.CommandCallback;

import static org.mockito.Matchers.any;
import static org.mockito.Mockito.doAnswer;
import static org.mockito.Mockito.mock;

public class CommandUtils<T> {

    public static Command mockSuccessCommand(Class<? extends Command> c) {
        return mockCommand(c, true, null);
    }

    public static Command mockSuccessCommand(Class<? extends Command> c, Object result) {
        return mockCommand(c, true, result);
    }

    public static Command mockFailCommand(Class<? extends Command> c) {
        return mockCommand(c, false, null);
    }

    private static Command mockCommand(Class<? extends Command> c, boolean success, Object result) {
        Command command = mock(c);
        try {
            doAnswer(invocation -> {
                CommandCallback callback = (CommandCallback) invocation.getArguments()[0];
                if (success) {
                    callback.onSuccess(result);
                } else {
                    callback.onFail(new Throwable("Command failed " + c.getSimpleName()));
                }
                return null;
            }).when(command).run(any(CommandCallback.class));
        } catch (Throwable throwable) {
            throwable.printStackTrace();
        }
        return command;
    }
}

MockCommandService.java

package com.ghm.di;

import com.ghm.service.api.CreatePSTNCallAction;
import com.ghm.service.util.SimpleActionServiceWrapper;

import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;

import io.techery.janet.ActionHolder;
import io.techery.janet.ActionService;
import io.techery.janet.Command;
import io.techery.janet.CommandUtils;

public class MockCommandService extends SimpleActionServiceWrapper {
    private Set<Class<? extends Command>> failCommand;
    private Map<Class<? extends Command>, Object> successCommand;

    MockCommandService(ActionService actionService) {
        super(actionService);
        failCommand = new HashSet<>();
        successCommand = new HashMap<>();
    }

    public void mockFail(Class<? extends Command> commandClass) {
        failCommand.add(commandClass);
        successCommand.remove(commandClass);
    }

    public void mockSuccess(Class<? extends Command> commandClass) {
        successCommand.put(commandClass, null);
        failCommand.remove(commandClass);
    }

    @Override protected <A> void onInterceptStart(ActionHolder<A> holder) {
        Class<? extends Command> commandClass = ((Command) holder.action()).getClass();
        if (successCommand.keySet().contains(commandClass)) {
            Object result = successCommand.get(commandClass);
            holder.newAction((A) CommandUtils.mockSuccessCommand(commandClass, result));
        }
        if (failCommand.contains(commandClass)) {
            holder.newAction((A) CommandUtils.mockFailCommand(commandClass));
        }
    }

    public void mockSuccess(Class<CreatePSTNCallAction> commandClass, Object result) {
            successCommand.put(commandClass, result);
    }
}