com.netflix.eureka2.registry.SourcedRegistryMockResource Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of eureka2-test-utils Show documentation
eureka2-test-utils
The newest version!
package com.netflix.eureka2.registry;

import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;
import java.util.Map.Entry;

import com.netflix.eureka2.interests.ChangeNotification;
import com.netflix.eureka2.interests.ChangeNotification.Kind;
import com.netflix.eureka2.interests.Interest;
import com.netflix.eureka2.interests.Interests;
import com.netflix.eureka2.interests.StreamStateNotification;
import com.netflix.eureka2.registry.instance.InstanceInfo;
import com.netflix.eureka2.testkit.data.builder.SampleInstanceInfo;
import org.junit.rules.ExternalResource;
import org.mockito.Mockito;
import org.mockito.invocation.InvocationOnMock;
import org.mockito.stubbing.Answer;
import rx.Observable;
import rx.Subscriber;
import rx.subjects.ReplaySubject;

import static org.mockito.Matchers.any;
import static org.mockito.Mockito.when;

/**
 * Provides {@link SourcedEurekaRegistry} mock with a set of helper methods to control it.
 *
 * @author Tomasz Bak
 */
public class SourcedRegistryMockResource extends ExternalResource {

    private final SourcedEurekaRegistry registry = Mockito.mock(SourcedEurekaRegistry.class);
    private volatile ReplaySubject> fullRegistrySubject = ReplaySubject.create();
    private final Map>> notificationSubjects = new HashMap<>();

    public SourcedEurekaRegistry registry() {
        return registry;
    }

    @Override
    protected void before() throws Throwable {
        when(registry.forInterest(any(Interest.class))).thenAnswer(new Answer>>() {
            @Override
            public Observable> answer(InvocationOnMock invocation) throws Throwable {
                return getInterestSubject((Interest) invocation.getArguments()[0]);
            }
        });
        fullRegistrySubject = ReplaySubject.create();
    }

    /**
     * Complete current notification subject, and create a new one. All interest subscribers will onComplete.
     */
    public void reset() {
        fullRegistrySubject.onCompleted();
        notificationSubjects.clear();

        Mockito.reset(registry);
        fullRegistrySubject = ReplaySubject.create();
    }

    public void batchStart() {
        batchStart(Interests.forFullRegistry());
    }

    public void batchStart(Interest interest) {
        getInterestSubject(interest).onNext(StreamStateNotification.bufferStartNotification(interest));
    }

    public void batchStartFor(InstanceInfo instanceInfo) {
        for (Entry>> entry : notificationSubjects.entrySet()) {
            if (entry.getKey().matches(instanceInfo)) {
                entry.getValue().onNext(StreamStateNotification.bufferStartNotification(entry.getKey()));
            }
        }
    }

    public void batchEnd() {
        batchEnd(Interests.forFullRegistry());
    }

    public void batchEnd(Interest interest) {
        getInterestSubject(interest).onNext(StreamStateNotification.bufferEndNotification(interest));
    }

    public void batchEndFor(InstanceInfo instanceInfo) {
        for (Entry>> entry : notificationSubjects.entrySet()) {
            if (entry.getKey().matches(instanceInfo)) {
                entry.getValue().onNext(StreamStateNotification.bufferEndNotification(entry.getKey()));
            }
        }
    }

    public void uploadToRegistry(InstanceInfo sample) {
        fullRegistrySubject.onNext(new ChangeNotification(Kind.Add, sample));
    }

    public void uploadBatchToRegistry(Interest interest, InstanceInfo sample) {
        batchStart(interest);
        fullRegistrySubject.onNext(new ChangeNotification(Kind.Add, sample));
        batchEnd(interest);
    }

    public void removeFromRegistry(InstanceInfo sample) {
        fullRegistrySubject.onNext(new ChangeNotification(Kind.Delete, sample));
    }

    /**
     * Upload a given number of instances to the registry, forged from the provided
     * template.
     */
    public void uploadClusterToRegistry(InstanceInfo template, int size) {
        Iterator instanceIt = SampleInstanceInfo.collectionOf(template.getApp(), template);
        for (int i = 0; i < size; i++) {
            uploadToRegistry(instanceIt.next());
        }
    }

    /**
     * Upload a given number of instances to the registry.
     */
    public String uploadClusterToRegistry(SampleInstanceInfo sample, int size) {
        String appName = null;
        for (InstanceInfo item : sample.clusterOf(size)) {
            appName = item.getApp();
            uploadToRegistry(item);
        }
        return appName;
    }

    public void uploadClusterBatchToRegistry(InstanceInfo template, int size) {
        Iterator instanceIt = SampleInstanceInfo.collectionOf(template.getApp(), template);

        InstanceInfo first = instanceIt.next();
        batchStartFor(first);
        uploadToRegistry(first);

        for (int i = 1; i < size; i++) {
            uploadToRegistry(instanceIt.next());
        }

        batchEndFor(first);
    }

    public String uploadClusterBatchToRegistry(SampleInstanceInfo sample, int size) {
        InstanceInfo first = null;
        for (InstanceInfo item : sample.clusterOf(size)) {
            if (first == null) {
                first = item;
                batchStartFor(item);
            }
            uploadToRegistry(item);
        }
        batchEndFor(first);
        return first.getApp();
    }

    private ReplaySubject> getInterestSubject(final Interest interest) {
        ReplaySubject> subject = notificationSubjects.get(interest);
        if (subject == null) {
            subject = ReplaySubject.create();
            final ReplaySubject> finalSubject = subject;
            fullRegistrySubject.subscribe(new Subscriber>() {
                @Override
                public void onCompleted() {
                    finalSubject.onCompleted();
                }

                @Override
                public void onError(Throwable e) {
                    finalSubject.onError(e);
                }

                @Override
                public void onNext(ChangeNotification notification) {
                    if (notification.isDataNotification() && interest.matches(notification.getData())) {
                        finalSubject.onNext(notification);
                    }
                }
            });
            notificationSubjects.put(interest, subject);
        }
        return subject;
    }
}