package fi.vincit.multiusertest.rule.expectation2;

import fi.vincit.multiusertest.rule.AuthorizationRule;
import fi.vincit.multiusertest.util.UserIdentifier;
import fi.vincit.multiusertest.util.UserIdentifiers;

import java.util.*;

public abstract class AbstractWhenThen<T extends TestExpectation> implements WhenThen<T> {

    private final Set<UserIdentifier> currentIdentifiers = new HashSet<>();
    private final Map<UserIdentifier, T> expectationsByIdentifier = new HashMap<>();

    private final UserIdentifier userIdentifier;
    private final AuthorizationRule authorizationRule;
    private T defaultExpectation;

    public AbstractWhenThen(UserIdentifier userIdentifier, AuthorizationRule authorizationRule) {
        this.userIdentifier = userIdentifier;
        this.authorizationRule = authorizationRule;
    }

    @Override
    public WhenThen<T> whenCalledWith(UserIdentifiers... userIdentifiers) {
        currentIdentifiers.clear();

        if (userIdentifiers.length == 0) {
            throw new IllegalArgumentException("At least one identifier must be defined");
        }

        for (UserIdentifiers identifiers : userIdentifiers) {
            identifiers.getIdentifiers().forEach(this::addCurrentUserIdentifiers);
        }
        return this;
    }

    @Override
    public WhenThen<T> whenCalledWith(UserIdentifier... userIdentifiers) {
        currentIdentifiers.clear();

        if (userIdentifiers.length == 0) {
            throw new IllegalArgumentException("At least one identifier must be defined");
        }

        addCurrentUserIdentifiers(userIdentifiers);

        return this;
    }

    private void addCurrentUserIdentifiers(UserIdentifier... userIdentifiers) {
        for (UserIdentifier identifier : userIdentifiers) {
            if (currentIdentifiers.contains(identifier)) {
                throw new IllegalStateException("User identifier " + identifier + " already set");
            }
            currentIdentifiers.add(identifier);
        }
    }

    @Override
    public WhenThen<T> then(final T testExpectation) {
        if (currentIdentifiers.isEmpty()) {
            throw new IllegalStateException("Call whenCalledWith before calling then method");
        }

        try {
            currentIdentifiers.forEach(identifier -> {
                if (expectationsByIdentifier.containsKey(identifier)) {
                    throw new IllegalStateException(
                            String.format("User identifier %s already has expectation",
                                    identifier.toString()
                            )
                    );
                } else {
                    expectationsByIdentifier.put(
                            identifier,
                            testExpectation
                    );
                }
            });
        } finally {
            currentIdentifiers.clear();
        }

        return this;
    }

    @Override
    public WhenThen<T> otherwise(T testExpectation) {
        this.defaultExpectation = testExpectation;
        return this;
    }

    @Override
    public WhenThen<T> byDefault(T testExpectation) {
        return otherwise(testExpectation);
    }

    @Override
    public void test() throws Throwable {
        T testExpectation = expectationsByIdentifier.computeIfAbsent(
                userIdentifier,
                this::getDefinedDefaultException
        );

        this.authorizationRule.markExpectationConstructed();
        this.test(testExpectation, userIdentifier);
    }

    protected abstract void test(T testExpectation, UserIdentifier userIdentifier) throws Throwable;

    protected abstract T getDefaultExpectation(UserIdentifier userIdentifier);

    protected T getDefinedDefaultException(UserIdentifier userIdentifier) {
        return Optional.ofNullable(defaultExpectation)
                .orElseGet(() -> getDefaultExpectation(userIdentifier));
    }

    /*
    For unit tests
    */

    Set<UserIdentifier> getCurrentIdentifiers() {
        return currentIdentifiers;
    }

    Map<UserIdentifier, T> getExpectationsByIdentifier() {
        return expectationsByIdentifier;
    }

}

@@ -2068,32 +2068,118 @@
 tExpectation) %7B%0A
+        Objects.requireNonNull(testExpectation, %22testExpectation must not be null%22);%0A%0A
         if (curr
@@ -3080,32 +3080,117 @@
 tExpectation) %7B%0A
+        Objects.requireNonNull(testExpectation, %22testExpectation must not be null%22);%0A
         this.def

Add null checks to AbstractWhenThen method arguments