package net.sf.twip.cdi;
   
   import java.util.Set;
   
   import javax.enterprise.context.spi.CreationalContext;
   import javax.enterprise.inject.spi.*;
   
   import net.sf.twip.TwipExtension;
   
  import org.jboss.weld.environment.se.Weld;
  import org.junit.runners.model.*;
  
  /**
   * Initializing CDI is expensive, so this {@link TwipExtension} is not auto-loaded... you'll have to
   * annotate your tests with <code>@TwipExtensions(CdiExtension.class)</code>.
   */
  public class CdiExtension extends TwipExtension {
  	private static final BeanManager beanManager = new Weld().initialize().getBeanManager();
  
  	private class CdiStatement<T> extends Statement {
  		private final Statement statement;
  		private final T instance;
  		private final CreationalContext<T> ctx;
  		private final InjectionTarget<T> target;
  
  		public CdiStatement(Statement statement, T instance) {
  			this.statement = statement;
  			this.instance = instance;
  
  			@SuppressWarnings("unchecked")
  			Class<T> testType = (Class<T>) instance.getClass();
  
  			Bean<T> bean = getUniqueBean(testType);
  			this.ctx = beanManager.createCreationalContext(bean);
  
  			AnnotatedType<T> type = beanManager.createAnnotatedType(testType);
  			this.target = beanManager.createInjectionTarget(type);
  		}
  
  		private Bean<T> getUniqueBean(Class<T> testType) {
  			Set<Bean<?>> beans = beanManager.getBeans(testType);
  			assert beans.size() == 1 : "there are " + beans.size() + " beans of " + testType;
  			@SuppressWarnings("unchecked")
  			Bean<T> bean = (Bean<T>) beans.iterator().next();
  			return bean;
  		}
  
  		@Override
  		public void evaluate() throws Throwable {
  			target.inject(instance, ctx);
  			target.postConstruct(instance);
  
  			try {
  				statement.evaluate();
  			} finally {
  				target.preDestroy(instance);
  				target.dispose(instance);
  				ctx.release();
  			}
  		}
  	}
  
  	@Override
  	public Statement wrapOuter(FrameworkMethod method, Object testInstance, Statement statement) {
  		return new CdiStatement<Object>(statement, testInstance);
  	}
  }