package test.org.mule.routing;

import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;

import org.mule.api.MuleException;
import org.mule.api.context.notification.ServerNotification;
import org.mule.module.client.MuleClient;
import org.mule.tck.functional.FunctionalTestNotificationListener;

public class MessageReceiver {

    final CountDownLatch latch;

    Object message = null;

    public MessageReceiver(MuleClient client) throws MuleException {
        latch = new CountDownLatch(1);
        client.getMuleContext().registerListener(
                new FunctionalTestNotificationListener() {
            public void onNotification(ServerNotification notification) {
                message = notification.getSource();
                latch.countDown();
            }
        });
    }

    public Object get(long timeout) throws InterruptedException {
        latch.await(timeout, TimeUnit.MILLISECONDS);
        if (message instanceof RuntimeException) {
            throw (RuntimeException) message;
        }
        return message;
    }

}
