package brooklyn.test;

import static org.testng.Assert.assertEquals;
import static org.testng.Assert.assertTrue;

import java.util.Collection;
import java.util.Map;

import brooklyn.config.ConfigKey;
import brooklyn.entity.Entity;
import brooklyn.entity.Group;
import brooklyn.event.AttributeSensor;

import com.google.common.base.Predicate;
import com.google.common.base.Predicates;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Maps;

public class EntityTestUtils {

	// TODO Delete methods from TestUtils, to just have them here (or switch so TestUtils delegates here,
	// and deprecate methods in TestUtils until deleted).
	
    public static <T> void assertAttributeEquals(Entity entity, AttributeSensor<T> attribute, T expected) {
        assertEquals(entity.getAttribute(attribute), expected);
    }
    
    public static <T> void assertConfigEquals(Entity entity, ConfigKey<T> configKey, T expected) {
        assertEquals(entity.getConfig(configKey), expected);
    }
    
    public static <T> void assertAttributeEqualsEventually(final Entity entity, final AttributeSensor<T> attribute, final T expected) {
        assertAttributeEqualsEventually(Maps.newLinkedHashMap(), entity, attribute, expected);
    }

    public static <T> void assertAttributeEqualsEventually(Map<?,?> flags, final Entity entity, final AttributeSensor<T> attribute, final T expected) {
        // Not using assertAttributeEventually(predicate) so get nicer error message
        Asserts.succeedsEventually((Map)flags, new Runnable() {
            @Override public void run() {
                assertAttributeEquals(entity, attribute, expected);
            }});
    }

    public static <T> void assertAttributeEventuallyNonNull(final Entity entity, final AttributeSensor<T> attribute) {
        assertAttributeEventuallyNonNull(Maps.newLinkedHashMap(), entity, attribute);
    }

    public static <T> void assertAttributeEventuallyNonNull(Map<?,?> flags, final Entity entity, final AttributeSensor<T> attribute) {
        assertAttributeEventually(flags, entity, attribute, Predicates.notNull());
    }

    public static <T> void assertAttributeEventually(final Entity entity, final AttributeSensor<T> attribute, Predicate<? super T> predicate) {
        assertAttributeEventually(ImmutableMap.of(), entity, attribute, predicate);
    }
    
    public static <T> void assertAttributeEventually(Map<?,?> flags, final Entity entity, final AttributeSensor<T> attribute, final Predicate<? super T> predicate) {
        Asserts.succeedsEventually((Map)flags, new Runnable() {
            @Override public void run() {
                T val = entity.getAttribute(attribute);
                assertTrue(predicate.apply(val), "val="+val);
            }});
    }

    public static <T extends Entity> void assertPredicateEventuallyTrue(final T entity, final Predicate<? super T> predicate) {
        assertPredicateEventuallyTrue(Maps.newLinkedHashMap(), entity, predicate);
    }

    public static <T extends Entity> void assertPredicateEventuallyTrue(Map<?,?> flags, final T entity, final Predicate<? super T> predicate) {
        Asserts.succeedsEventually((Map)flags, new Runnable() {
                @Override public void run() {
                    assertTrue(predicate.apply(entity));
                }});
    }

    public static <T> void assertAttributeEqualsContinually(final Entity entity, final AttributeSensor<T> attribute, final T expected) {
        assertAttributeEqualsContinually(Maps.newLinkedHashMap(), entity, attribute, expected);
    }
    
    public static <T> void assertAttributeEqualsContinually(Map<?,?> flags, final Entity entity, final AttributeSensor<T> attribute, final T expected) {
        Asserts.succeedsContinually(flags, new Runnable() {
                @Override public void run() {
                    assertAttributeEquals(entity, attribute, expected);
                }});
    }

    public static void assertGroupSizeEqualsEventually(final Group group, int expected) {
        assertGroupSizeEqualsEventually(ImmutableMap.of(), group, expected);
    }
    
    public static void assertGroupSizeEqualsEventually(Map<?,?> flags, final Group group, final int expected) {
        Asserts.succeedsEventually((Map)flags, new Runnable() {
            @Override public void run() {
                Collection<Entity> members = group.getMembers();
                assertEquals(members.size(), expected, "members="+members);
            }});
    }
}

@@ -158,16 +158,68 @@
 til.Map;
+%0Aimport java.util.concurrent.atomic.AtomicReference;
 %0A%0Aimport
@@ -1750,36 +1750,33 @@
 blic static %3CT%3E 
-void
+T
  assertAttribute
@@ -1852,32 +1852,39 @@
 bute) %7B%0A        
+return 
 assertAttributeE
@@ -1966,36 +1966,33 @@
 blic static %3CT%3E 
-void
+T
  assertAttribute
@@ -2084,32 +2084,39 @@
 bute) %7B%0A        
+return 
 assertAttributeE
@@ -2195,36 +2195,33 @@
 blic static %3CT%3E 
-void
+T
  assertAttribute
@@ -2322,32 +2322,39 @@
 cate) %7B%0A        
+return 
 assertAttributeE
@@ -2438,36 +2438,33 @@
 blic static %3CT%3E 
-void
+T
  assertAttribute
@@ -2579,32 +2579,100 @@
 T%3E predicate) %7B%0A
+        final AtomicReference%3CT%3E result = new AtomicReference%3CT%3E();%0A
         Asserts.
@@ -2883,34 +2883,96 @@
 l);%0A            
-%7D%7D
+    result.set(val);%0A            %7D%7D);%0A        return result.get(
 );%0A    %7D%0A%0A    pu

return val