RunAs

RunAsSupportTestExecutionListener.java

public class RunAsSupportTestExecutionListener extends AbstractTestExecutionListener {
    private static final ThreadLocal<User> savedAuthentication = new ThreadLocal<User>();

	@Override
	public void prepareTestInstance(TestContext testContext) throws Exception {
		super.prepareTestInstance(testContext);
		try {
			testContext.getApplicationContext().getBean(CpsModuleSecuritySupport.class);
		} catch (NoSuchBeanDefinitionException ignored) {
			CpsModuleSupport moduleSupport = mock(CpsModuleSupport.class);
			ConfigurableListableBeanFactory beanFactory = (ConfigurableListableBeanFactory)testContext.getApplicationContext().getAutowireCapableBeanFactory();
			beanFactory.registerSingleton("cpsModuleSupport", moduleSupport);
		}
	}

	@Override
    public void beforeTestMethod(TestContext testContext) throws Exception {
        super.beforeTestMethod(testContext);
        final RunAs runAs = testContext.getTestMethod().getAnnotation(RunAs.class);
        if (runAs != null) {
            final String userName = runAs.value();
            loginAsUser(userName, testContext.getApplicationContext());
        }
    }

    @Override
    public void afterTestMethod(TestContext testContext) throws Exception {
        super.afterTestMethod(testContext);
        final RunAs runAs = testContext.getTestMethod().getAnnotation(RunAs.class);
        if (runAs != null) {
            logoutUser(testContext.getApplicationContext());
        }
    }

    public static void loginAsUser(String userName, ApplicationContext appCtx) {
		User user = createUser(userName);
		loginAsUser(user, appCtx);
	}

	public static void loginAsUser(User user, ApplicationContext appCtx) {
		CpsModuleSupport cpsModuleSupport = appCtx.getBean(CpsModuleSupport.class);
		CpsModuleSecuritySupport securitySupport = mock(CpsModuleSecuritySupport.class);
		when(cpsModuleSupport.getSecuritySupport()).thenReturn(securitySupport);
		when(securitySupport.getAuthenticatedUser(any(HttpServletRequest.class), any(HttpServletResponse.class))).thenReturn(user);
		UserProviderFilter.setLoggedInUser(user);
	}

	private static void logoutUser(ApplicationContext appCtx) {
		CpsModuleSupport cpsModuleSupport = appCtx.getBean(CpsModuleSupport.class);
		CpsModuleSecuritySupport securitySupport = mock(CpsModuleSecuritySupport.class);
		when(cpsModuleSupport.getSecuritySupport()).thenReturn(securitySupport);
		when(securitySupport.getAuthenticatedUser(any(HttpServletRequest.class), any(HttpServletResponse.class))).thenReturn(null);
		UserProviderFilter.setLoggedInUser(null);
    }

    public static void executeInContextOf(String userName, ApplicationContext appCtx, ExecutionClosure closure) {
        try {
            pushUser(createUser(userName), appCtx);
            closure.execute();
        } finally {
            popUser(appCtx);
        }
    }

	private static User createUser(String userName) {
		return new User(userName);
	}

	private static void pushUser(User user, ApplicationContext appCtx) {
		savedAuthentication.set(UserProviderFilter.getLoggedInUser());
		loginAsUser(user, appCtx);
	}

	private static void popUser(ApplicationContext appCtx) {
		User user = savedAuthentication.get();
		savedAuthentication.remove();
		if (user == null) {
			logoutUser(appCtx);
		} else {
			loginAsUser(user, appCtx);
		}
	}

    public interface ExecutionClosure {

        void execute();

    }

}