text:
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

package com.azure.cosmos.implementation.directconnectivity;

import com.azure.cosmos.ConsistencyLevel;
import com.azure.cosmos.implementation.NotFoundException;
import com.azure.cosmos.implementation.RequestRateTooLargeException;
import com.azure.cosmos.implementation.Configs;
import com.azure.cosmos.implementation.DocumentServiceRequestContext;
import com.azure.cosmos.implementation.FailureValidator;
import com.azure.cosmos.implementation.HttpConstants;
import com.azure.cosmos.implementation.IAuthorizationTokenProvider;
import com.azure.cosmos.implementation.ISessionContainer;
import com.azure.cosmos.implementation.ISessionToken;
import com.azure.cosmos.implementation.OperationType;
import com.azure.cosmos.implementation.PartitionKeyRange;
import com.azure.cosmos.implementation.RequestChargeTracker;
import com.azure.cosmos.implementation.ResourceType;
import com.azure.cosmos.implementation.RxDocumentServiceRequest;
import com.azure.cosmos.implementation.StoreResponseBuilder;
import com.azure.cosmos.implementation.Utils;
import com.azure.cosmos.implementation.VectorSessionToken;
import com.azure.cosmos.implementation.guava25.collect.ImmutableList;
import io.reactivex.subscribers.TestSubscriber;
import org.mockito.Mockito;
import org.testng.annotations.DataProvider;
import org.testng.annotations.Test;
import reactor.core.publisher.Mono;

import java.math.BigDecimal;
import java.math.RoundingMode;
import java.util.List;
import java.util.concurrent.TimeUnit;

import static com.azure.cosmos.implementation.Utils.ValueHolder;
import static org.assertj.core.api.Assertions.assertThat;
import static com.azure.cosmos.implementation.TestUtils.*;

public class ConsistencyReaderTest {
    private final Configs configs = new Configs();
    private static final int TIMEOUT = 30000;
    @DataProvider(name = "deduceReadModeArgProvider")
    public Object[][] deduceReadModeArgProvider() {
        return new Object[][]{
                // account consistency, request consistency, expected readmode, expected consistency to use, whether use session
                {  ConsistencyLevel.STRONG, null, ReadMode.Strong, ConsistencyLevel.STRONG, false},
                {  ConsistencyLevel.STRONG, ConsistencyLevel.EVENTUAL, ReadMode.Any, ConsistencyLevel.EVENTUAL, false},
                {  ConsistencyLevel.STRONG, ConsistencyLevel.SESSION, ReadMode.Any, ConsistencyLevel.SESSION, true},
                {  ConsistencyLevel.SESSION, ConsistencyLevel.EVENTUAL, ReadMode.Any, ConsistencyLevel.EVENTUAL, false},
                {  ConsistencyLevel.SESSION, ConsistencyLevel.SESSION, ReadMode.Any, ConsistencyLevel.SESSION, true},
                {  ConsistencyLevel.SESSION, ConsistencyLevel.EVENTUAL, ReadMode.Any, ConsistencyLevel.EVENTUAL, false},
                {  ConsistencyLevel.SESSION, null, ReadMode.Any, ConsistencyLevel.SESSION, true},
                {  ConsistencyLevel.EVENTUAL, ConsistencyLevel.EVENTUAL, ReadMode.Any, ConsistencyLevel.EVENTUAL, false},
                {  ConsistencyLevel.EVENTUAL, null, ReadMode.Any, ConsistencyLevel.EVENTUAL, false},
        };
    }

    @Test(groups = "unit", dataProvider = "deduceReadModeArgProvider")
    public void deduceReadMode(ConsistencyLevel accountConsistencyLevel, ConsistencyLevel requestConsistency, ReadMode expectedReadMode,
                               ConsistencyLevel expectedConsistencyToUse, boolean expectedToUseSession) {
        AddressSelector addressSelector = Mockito.mock(AddressSelector.class);
        ISessionContainer sessionContainer = Mockito.mock(ISessionContainer.class);
        TransportClient transportClient = Mockito.mock(TransportClient.class);
        GatewayServiceConfiguratorReaderMock gatewayServiceConfigurationReaderWrapper = GatewayServiceConfiguratorReaderMock.from(accountConsistencyLevel);
        IAuthorizationTokenProvider authorizationTokenProvider = Mockito.mock(IAuthorizationTokenProvider.class);
        ConsistencyReader consistencyReader = new ConsistencyReader(mockDiagnosticsClientContext(),
                                                                    configs,
                                                                    addressSelector,
                                                                    sessionContainer,
                                                                    transportClient,
                                                                    gatewayServiceConfigurationReaderWrapper.gatewayServiceConfigurationReader,
                                                                    authorizationTokenProvider);

        RxDocumentServiceRequest request = RxDocumentServiceRequest.createFromName(mockDiagnosticsClientContext(),
                OperationType.Read, "/dbs/db/colls/col/docs/docId", ResourceType.Document);
        if (requestConsistency != null) {
            request.getHeaders().put(HttpConstants.HttpHeaders.CONSISTENCY_LEVEL, requestConsistency.toString());
        }

        ValueHolder<ConsistencyLevel> consistencyLevel = ValueHolder.initialize(null);
        ValueHolder<Boolean> useSession = ValueHolder.initialize(null);

        ReadMode readMode = consistencyReader.deduceReadMode(request, consistencyLevel, useSession);

        assertThat(readMode).isEqualTo(expectedReadMode);
        assertThat(consistencyLevel.v).isEqualTo(expectedConsistencyToUse);
        assertThat(useSession.v).isEqualTo(expectedToUseSession);
    }

    @DataProvider(name = "getMaxReplicaSetSizeArgProvider")
    public Object[][] getMaxReplicaSetSizeArgProvider() {
        return new Object[][]{
                // system max replica count, system min replica count, user max replica count, user min replica, is reading from master operation
                {  4, 3, 4, 3, false },
                {  4, 3, 4, 3, true },

                {  4, 3, 3, 2, false },
                {  4, 3, 3, 2, true }
        };
    }

    @Test(groups = "unit", dataProvider = "getMaxReplicaSetSizeArgProvider")
    public void replicaSizes(int systemMaxReplicaCount,
                                     int systemMinReplicaCount,
                                     int userMaxReplicaCount,
                                     int userMinReplicaCount,
                                     boolean isReadingFromMasterOperation) {
        AddressSelector addressSelector = Mockito.mock(AddressSelector.class);
        ISessionContainer sessionContainer = Mockito.mock(ISessionContainer.class);
        TransportClient transportClient = Mockito.mock(TransportClient.class);
        GatewayServiceConfiguratorReaderMock gatewayServiceConfigurationReaderWrapper = GatewayServiceConfiguratorReaderMock.from(ConsistencyLevel.STRONG,
                                                                                systemMaxReplicaCount,
                                                                                systemMinReplicaCount,
                                                                                userMaxReplicaCount,
                                                                                userMinReplicaCount);
        IAuthorizationTokenProvider authorizationTokenProvider = Mockito.mock(IAuthorizationTokenProvider.class);
        ConsistencyReader consistencyReader = new ConsistencyReader(mockDiagnosticsClientContext(),
                                                                    configs,
                                                                    addressSelector,
                                                                    sessionContainer,
                                                                    transportClient,
                                                                    gatewayServiceConfigurationReaderWrapper.gatewayServiceConfigurationReader,
                                                                    authorizationTokenProvider);

        RxDocumentServiceRequest request;
        if (isReadingFromMasterOperation) {
            request = RxDocumentServiceRequest.createFromName(mockDiagnosticsClientContext(),
                    OperationType.ReadFeed, "/dbs/db/colls/col", ResourceType.DocumentCollection);
        } else {
            request = RxDocumentServiceRequest.createFromName(mockDiagnosticsClientContext(),
                    OperationType.Read, "/dbs/db/colls/col/docs/docId", ResourceType.Document);
        }

        assertThat(consistencyReader.getMaxReplicaSetSize(request)).isEqualTo(isReadingFromMasterOperation? systemMaxReplicaCount : userMaxReplicaCount);
        assertThat(consistencyReader.getMinReplicaSetSize(request)).isEqualTo(isReadingFromMasterOperation? systemMinReplicaCount : userMinReplicaCount);
    }

    @Test(groups = "unit")
    public void readAny() {
        List<Uri> secondaries = ImmutableList.of(Uri.create("secondary1"), Uri.create("secondary2"), Uri.create("secondary3"));
        Uri primaryAddress = Uri.create("primary");
        AddressSelectorWrapper addressSelectorWrapper = AddressSelectorWrapper.Builder.Simple.create()
                .withPrimary(primaryAddress)
                .withSecondary(secondaries)
                .build();

        StoreResponse primaryResponse = StoreResponseBuilder.create()
                .withLSN(54)
                .withLocalLSN(18)
                .withRequestCharge(1.1)
                .build();
        StoreResponse secondaryResponse1 = StoreResponseBuilder.create()
                .withLSN(53)
                .withLocalLSN(17)
                .withRequestCharge(1.1)
                .build();
        StoreResponse secondaryResponse2 = StoreResponseBuilder.create()
                .withLSN(52)
                .withLocalLSN(16)
                .withRequestCharge(1.1)
                .build();
        StoreResponse secondaryResponse3 = StoreResponseBuilder.create()
                .withLSN(51)
                .withLocalLSN(15)
                .withRequestCharge(1.1)
                .build();
        TransportClientWrapper transportClientWrapper = TransportClientWrapper.Builder.uriToResultBuilder()
                .storeResponseOn(primaryAddress, OperationType.Read, ResourceType.Document, primaryResponse, true)
                .storeResponseOn(secondaries.get(0), OperationType.Read, ResourceType.Document, secondaryResponse1, true)
                .storeResponseOn(secondaries.get(1), OperationType.Read, ResourceType.Document, secondaryResponse2, true)
                .storeResponseOn(secondaries.get(2), OperationType.Read, ResourceType.Document, secondaryResponse3, true)
                .build();

        GatewayServiceConfiguratorReaderMock gatewayServiceConfigurationReaderWrapper = GatewayServiceConfiguratorReaderMock.from(ConsistencyLevel.STRONG);
        ISessionContainer sessionContainer = Mockito.mock(ISessionContainer.class);

        IAuthorizationTokenProvider authorizationTokenProvider = Mockito.mock(IAuthorizationTokenProvider.class);
        ConsistencyReaderUnderTest consistencyReader = new ConsistencyReaderUnderTest(addressSelectorWrapper.addressSelector,
                                                                                sessionContainer,
                                                                                transportClientWrapper.transportClient,
                                                                                gatewayServiceConfigurationReaderWrapper.gatewayServiceConfigurationReader,
                                                                                authorizationTokenProvider);

        RxDocumentServiceRequest request = RxDocumentServiceRequest.createFromName(mockDiagnosticsClientContext(),
                OperationType.Read, "/dbs/db/colls/col/docs/docId", ResourceType.Document);
        request.getHeaders().put(HttpConstants.HttpHeaders.CONSISTENCY_LEVEL, ConsistencyLevel.EVENTUAL.toString());

        TimeoutHelper timeout = Mockito.mock(TimeoutHelper.class);
        boolean forceRefresh = false;
        boolean isInRetry = false;
        Mono<StoreResponse> storeResponseSingle = consistencyReader.readAsync(request, timeout, isInRetry, forceRefresh);

        StoreResponseValidator validator = StoreResponseValidator.create()
                .withBELSNGreaterThanOrEqualTo(51)
                .withRequestCharge(1.1)
                .in(primaryResponse, secondaryResponse1, secondaryResponse2, secondaryResponse3)
                .build();
        validateSuccess(storeResponseSingle, validator);

        Mockito.verifyZeroInteractions(consistencyReader.getSpyQuorumReader());

        Mockito.verify(consistencyReader.getSpyStoreReader(), Mockito.times(1))
                .readMultipleReplicaAsync(Mockito.any(RxDocumentServiceRequest.class),
                                          Mockito.anyBoolean(),
                                          Mockito.anyInt(),
                                          Mockito.anyBoolean(),
                                          Mockito.anyBoolean(),
                                          Mockito.any(),
                                          Mockito.anyBoolean(),
                                          Mockito.anyBoolean());

        transportClientWrapper.validate()
                .verifyNumberOfInvocations(1);

        addressSelectorWrapper.validate()
                .verifyTotalInvocations(1)
                .verifyNumberOfForceCachRefresh(0)
                .verifyVesolvePrimaryUriAsyncCount(0)
                .verifyResolveAllUriAsync(1);
    }

    @Test(groups = "unit")
    public void readSessionConsistency_SomeReplicasLagBehindAndReturningResponseWithLowerLSN_FindAnotherReplica() {
        long slowReplicaLSN = 651176;
        String partitionKeyRangeId = "1";
        long fasterReplicaLSN = 651177;

        List<Uri> secondaries = ImmutableList.of(Uri.create("secondary1"), Uri.create("secondary2"), Uri.create("secondary3"));
        Uri primaryAddress = Uri.create("primary");
        AddressSelectorWrapper addressSelectorWrapper = AddressSelectorWrapper.Builder.Simple.create()
                .withPrimary(primaryAddress)
                .withSecondary(secondaries)
                .build();

        StoreResponse primaryResponse = StoreResponseBuilder.create()
                .withSessionToken(partitionKeyRangeId + ":-1#" + slowReplicaLSN)
                .withLSN(slowReplicaLSN)
                .withLocalLSN(slowReplicaLSN)
                .withQuorumAckecdLsn(slowReplicaLSN)
                .withQuorumAckecdLocalLsn(slowReplicaLSN)
                .withGlobalCommittedLsn(-1)
                .withItemLocalLSN(slowReplicaLSN)
                .withRequestCharge(1.1)
                .build();
        StoreResponse secondaryResponse1 = StoreResponseBuilder.create()
                .withSessionToken(partitionKeyRangeId + ":-1#" + slowReplicaLSN)
                .withLSN(slowReplicaLSN)
                .withLocalLSN(slowReplicaLSN)
                .withQuorumAckecdLsn(slowReplicaLSN)
                .withQuorumAckecdLocalLsn(slowReplicaLSN)
                .withGlobalCommittedLsn(-1)
                .withItemLocalLSN(slowReplicaLSN)
                .withRequestCharge(1.1)
                .build();
        StoreResponse secondaryResponse2 = StoreResponseBuilder.create()
                .withSessionToken(partitionKeyRangeId + ":-1#" + fasterReplicaLSN)
                .withLSN(fasterReplicaLSN)
                .withLocalLSN(fasterReplicaLSN)
                .withQuorumAckecdLsn(fasterReplicaLSN)
                .withQuorumAckecdLocalLsn(fasterReplicaLSN)
                .withGlobalCommittedLsn(-1)
                .withItemLocalLSN(fasterReplicaLSN)
                .withRequestCharge(1.1)
                .build();
        StoreResponse secondaryResponse3 = StoreResponseBuilder.create()
                .withSessionToken(partitionKeyRangeId + ":-1#" + slowReplicaLSN)
                .withLSN(slowReplicaLSN)
                .withLocalLSN(slowReplicaLSN)
                .withQuorumAckecdLsn(slowReplicaLSN)
                .withQuorumAckecdLocalLsn(slowReplicaLSN)
                .withGlobalCommittedLsn(-1)
                .withItemLocalLSN(slowReplicaLSN)
                .withRequestCharge(1.1)
                .build();
        TransportClientWrapper transportClientWrapper = TransportClientWrapper.Builder.uriToResultBuilder()
                .storeResponseOn(primaryAddress, OperationType.Read, ResourceType.Document, primaryResponse, true)
                .storeResponseOn(secondaries.get(0), OperationType.Read, ResourceType.Document, secondaryResponse1, true)
                .storeResponseOn(secondaries.get(1), OperationType.Read, ResourceType.Document, secondaryResponse2, true)
                .storeResponseOn(secondaries.get(2), OperationType.Read, ResourceType.Document, secondaryResponse3, true)
                .build();

        GatewayServiceConfiguratorReaderMock gatewayServiceConfigurationReaderWrapper = GatewayServiceConfiguratorReaderMock.from(ConsistencyLevel.STRONG);
        ISessionContainer sessionContainer = Mockito.mock(ISessionContainer.class);

        IAuthorizationTokenProvider authorizationTokenProvider = Mockito.mock(IAuthorizationTokenProvider.class);
        ConsistencyReaderUnderTest consistencyReader = new ConsistencyReaderUnderTest(addressSelectorWrapper.addressSelector,
                                                                                sessionContainer,
                                                                                transportClientWrapper.transportClient,
                                                                                gatewayServiceConfigurationReaderWrapper.gatewayServiceConfigurationReader,
                                                                                authorizationTokenProvider);

        RxDocumentServiceRequest request = RxDocumentServiceRequest.createFromName(mockDiagnosticsClientContext(),
                OperationType.Read, "/dbs/db/colls/col/docs/docId", ResourceType.Document);
        request.getHeaders().put(HttpConstants.HttpHeaders.CONSISTENCY_LEVEL, ConsistencyLevel.SESSION.toString());
        request.requestContext = new DocumentServiceRequestContext();
        Utils.ValueHolder<ISessionToken> sessionToken = Utils.ValueHolder.initialize(null);
        assertThat(VectorSessionToken.tryCreate("-1#" + fasterReplicaLSN , sessionToken)).isTrue();
        request.requestContext.timeoutHelper = Mockito.mock(TimeoutHelper.class);
        request.requestContext.resolvedPartitionKeyRange = partitionKeyRangeWithId(partitionKeyRangeId);
        request.requestContext.requestChargeTracker = new RequestChargeTracker();

        Mockito.doReturn(sessionToken.v).when(sessionContainer).resolvePartitionLocalSessionToken(Mockito.eq(request), Mockito.anyString());

        TimeoutHelper timeout = Mockito.mock(TimeoutHelper.class);
        boolean forceRefresh = false;
        boolean isInRetry = false;
        Mono<StoreResponse> storeResponseSingle = consistencyReader.readAsync(request, timeout, isInRetry, forceRefresh);

        StoreResponseValidator validator = StoreResponseValidator.create()
                .withBELSN(fasterReplicaLSN)
                .withRequestChargeGreaterThanOrEqualTo(1.1)
                .in(primaryResponse, secondaryResponse1, secondaryResponse2, secondaryResponse3)
                .build();
        validateSuccess(storeResponseSingle, validator);

        Mockito.verifyZeroInteractions(consistencyReader.getSpyQuorumReader());

        Mockito.verify(consistencyReader.getSpyStoreReader(), Mockito.times(1))
                .readMultipleReplicaAsync(Mockito.any(RxDocumentServiceRequest.class),
                                          Mockito.anyBoolean(),
                                          Mockito.anyInt(),
                                          Mockito.anyBoolean(),
                                          Mockito.anyBoolean(),
                                          Mockito.any(),
                                          Mockito.anyBoolean(),
                                          Mockito.anyBoolean());

        assertThat(transportClientWrapper.validate()
                .getNumberOfInvocations())
                .isGreaterThanOrEqualTo(1)
                .isLessThanOrEqualTo(4);

        addressSelectorWrapper.validate()
                .verifyTotalInvocations(1)
                .verifyNumberOfForceCachRefresh(0)
                .verifyVesolvePrimaryUriAsyncCount(0)
                .verifyResolveAllUriAsync(1);
    }

    /**
     * reading in session consistency, if the requested session token cannot be supported by some replicas
     * tries others till we find a replica which can support the given session token
     */
    @Test(groups = "unit")
    public void sessionNotAvailableFromSomeReplicasThrowingNotFound_FindReplicaSatisfyingRequestedSession() {
        long slowReplicaLSN = 651175;
        long globalCommittedLsn = 651174;

        long fasterReplicaLSN = 651176;
        String partitionKeyRangeId = "1";

        NotFoundException foundException = new NotFoundException();
        foundException.getResponseHeaders().put(HttpConstants.HttpHeaders.SESSION_TOKEN, partitionKeyRangeId + ":-1#" + slowReplicaLSN);
        foundException.getResponseHeaders().put(WFConstants.BackendHeaders.LSN, Long.toString(slowReplicaLSN));
        foundException.getResponseHeaders().put(WFConstants.BackendHeaders.LOCAL_LSN, Long.toString(slowReplicaLSN));
        foundException.getResponseHeaders().put(WFConstants.BackendHeaders.GLOBAL_COMMITTED_LSN, Long.toString(globalCommittedLsn));

        StoreResponse storeResponse = StoreResponseBuilder.create()
                .withSessionToken(partitionKeyRangeId + ":-1#" + fasterReplicaLSN)
                .withLSN(fasterReplicaLSN)
                .withLocalLSN(fasterReplicaLSN)
                .withQuorumAckecdLsn(fasterReplicaLSN)
                .withQuorumAckecdLocalLsn(fasterReplicaLSN)
                .withGlobalCommittedLsn(-1)
                .withItemLocalLSN(fasterReplicaLSN)
                .withRequestCharge(1.1)
                .build();

        TransportClientWrapper transportClientWrapper = new TransportClientWrapper.Builder.ReplicaResponseBuilder
                .SequentialBuilder()
                .then(foundException) // 1st replica read returns not found
                .then(foundException) // 2nd replica read returns not found
                .then(foundException) // 3rd replica read returns not found
                .then(storeResponse)  // 4th replica read returns storeResponse satisfying requested session token
                .build();

        Uri primaryUri = Uri.create("primary");
        Uri secondaryUri1 = Uri.create("secondary1");
        Uri secondaryUri2 = Uri.create("secondary2");
        Uri secondaryUri3 = Uri.create("secondary3");

        AddressSelectorWrapper addressSelectorWrapper = AddressSelectorWrapper.Builder.Simple.create()
                .withPrimary(primaryUri)
                .withSecondary(ImmutableList.of(secondaryUri1, secondaryUri2, secondaryUri3))
                .build();
        ISessionContainer sessionContainer = Mockito.mock(ISessionContainer.class);

        Configs configs = new Configs();
        GatewayServiceConfiguratorReaderMock gatewayServiceConfigurationReaderWrapper = GatewayServiceConfiguratorReaderMock.from(ConsistencyLevel.STRONG,
                                                                                4,
                                                                                3,
                                                                                4,
                                                                                3);

        IAuthorizationTokenProvider authTokenProvider = Mockito.mock(IAuthorizationTokenProvider.class);
        ConsistencyReader consistencyReader = new ConsistencyReader(mockDiagnosticsClientContext(),
                                                                    configs,
                                                                    addressSelectorWrapper.addressSelector,
                                                                    sessionContainer,
                                                                    transportClientWrapper.transportClient,
                                                                    gatewayServiceConfigurationReaderWrapper.gatewayServiceConfigurationReader,
                                                                    authTokenProvider);

        TimeoutHelper timeoutHelper = Mockito.mock(TimeoutHelper.class);
        RxDocumentServiceRequest dsr = RxDocumentServiceRequest.createFromName(mockDiagnosticsClientContext(),
                OperationType.Read, "/dbs/db/colls/col/docs/docId", ResourceType.Document);
        dsr.getHeaders().put(HttpConstants.HttpHeaders.CONSISTENCY_LEVEL, ConsistencyLevel.SESSION.toString());
        dsr.requestContext = new DocumentServiceRequestContext();
        Utils.ValueHolder<ISessionToken> sessionToken = Utils.ValueHolder.initialize(null);
        assertThat(VectorSessionToken.tryCreate("-1#" + fasterReplicaLSN , sessionToken)).isTrue();
        dsr.requestContext.timeoutHelper = timeoutHelper;
        dsr.requestContext.resolvedPartitionKeyRange = partitionKeyRangeWithId(partitionKeyRangeId);
        dsr.requestContext.requestChargeTracker = new RequestChargeTracker();

        Mockito.doReturn(sessionToken.v).when(sessionContainer).resolvePartitionLocalSessionToken(Mockito.eq(dsr), Mockito.anyString());

        Mono<StoreResponse> storeResponseSingle = consistencyReader.readAsync(dsr, timeoutHelper, false, false);

        StoreResponseValidator validator = StoreResponseValidator.create().isSameAs(storeResponse).isSameAs(storeResponse).build();
        validateSuccess(storeResponseSingle, validator);
        transportClientWrapper.verifyNumberOfInvocations(4);
    }

    /**
     * reading in session consistency, if the requested session token cannot be supported by some replicas
     * tries others till we find a replica which can support the given session token
     */
    @Test(groups = "unit")
    public void sessionNotAvailableFromAllReplicasThrowingNotFound_FindReplicaSatisfyingRequestedSessionOnRetry() {
        long slowReplicaLSN = 651175;
        long globalCommittedLsn = 651174;

        long fasterReplicaLSN = 651176;
        String partitionKeyRangeId = "1";

        NotFoundException notFoundException = new NotFoundException();
        notFoundException.getResponseHeaders().put(
            HttpConstants.HttpHeaders.SESSION_TOKEN,
            partitionKeyRangeId + ":-1#" + slowReplicaLSN);
        notFoundException.getResponseHeaders().put(
            WFConstants.BackendHeaders.LSN,
            Long.toString(slowReplicaLSN));
        notFoundException.getResponseHeaders().put(
            WFConstants.BackendHeaders.LOCAL_LSN,
            Long.toString(slowReplicaLSN));
        notFoundException.getResponseHeaders().put(
            WFConstants.BackendHeaders.GLOBAL_COMMITTED_LSN,
            Long.toString(globalCommittedLsn));

        StoreResponse storeResponse =
            StoreResponseBuilder.create()
                                .withSessionToken(partitionKeyRangeId + ":-1#" + fasterReplicaLSN)
                                .withLSN(fasterReplicaLSN)
                                .withLocalLSN(fasterReplicaLSN)
                                .withQuorumAckecdLsn(fasterReplicaLSN)
                                .withQuorumAckecdLocalLsn(fasterReplicaLSN)
                                .withGlobalCommittedLsn(-1)
                                .withItemLocalLSN(fasterReplicaLSN)
                                .withRequestCharge(1.1)
                                .build();

        TransportClientWrapper transportClientWrapper = new TransportClientWrapper.Builder.ReplicaResponseBuilder
            .SequentialBuilder()
            .then(notFoundException) // 1st replica read returns not found
            .then(notFoundException) // 2nd replica read returns not found
            .then(notFoundException) // 3rd replica read returns not found
            .then(notFoundException) // 4th replica read returns not found
            .then(notFoundException) // 1st replica read returns not found - 1st retry
            .then(notFoundException) // 2nd replica read returns not found - 1st retry
            .then(notFoundException) // 3rd replica read returns not found - 1st retry
            .then(notFoundException) // 4th replica read returns not found - 1st retry
            .then(notFoundException) // 1st replica read returns not found - 2nd retry
            .then(notFoundException) // 2nd replica read returns not found - 2nd retry
            .then(notFoundException) // 3rd replica read returns not found - 2nd retry
            .then(storeResponse)  // 4th replica returns storeResponse satisfying requested session token on 2nd retry
            .build();

        Uri primaryUri = Uri.create("primary");
        Uri secondaryUri1 = Uri.create("secondary1");
        Uri secondaryUri2 = Uri.create("secondary2");
        Uri secondaryUri3 = Uri.create("secondary3");
        ImmutableList<Uri> secondaryUris = ImmutableList.of(secondaryUri1, secondaryUri2,
            secondaryUri3);

        AddressSelectorWrapper addressSelectorWrapper =
            AddressSelectorWrapper.Builder.Simple.create()
                                                 .withPrimary(primaryUri)
                                                 .withSecondary(secondaryUris)
                                                 .build();
        ISessionContainer sessionContainer = Mockito.mock(ISessionContainer.class);

        Configs configs = new Configs();
        GatewayServiceConfiguratorReaderMock gatewayServiceConfigurationReaderWrapper =
            GatewayServiceConfiguratorReaderMock.from(ConsistencyLevel.STRONG,
            4,
            3,
            4,
            3);

        IAuthorizationTokenProvider authTokenProvider = Mockito.mock(IAuthorizationTokenProvider.class);
        ConsistencyReader consistencyReader = new ConsistencyReader(mockDiagnosticsClientContext(),
            configs,
            addressSelectorWrapper.addressSelector,
            sessionContainer,
            transportClientWrapper.transportClient,
            gatewayServiceConfigurationReaderWrapper.gatewayServiceConfigurationReader,
            authTokenProvider);

        TimeoutHelper timeoutHelper = Mockito.mock(TimeoutHelper.class);
        RxDocumentServiceRequest dsr = RxDocumentServiceRequest.createFromName(mockDiagnosticsClientContext(),
            OperationType.Read, "/dbs/db/colls/col/docs/docId", ResourceType.Document);
        dsr.getHeaders().put(HttpConstants.HttpHeaders.CONSISTENCY_LEVEL, ConsistencyLevel.SESSION.toString());
        dsr.requestContext = new DocumentServiceRequestContext();
        Utils.ValueHolder<ISessionToken> sessionToken = Utils.ValueHolder.initialize(null);
        assertThat(VectorSessionToken.tryCreate("-1#" + fasterReplicaLSN, sessionToken)).isTrue();
        dsr.requestContext.timeoutHelper = timeoutHelper;
        dsr.requestContext.resolvedPartitionKeyRange = partitionKeyRangeWithId(partitionKeyRangeId);
        dsr.requestContext.requestChargeTracker = new RequestChargeTracker();

        Mockito
            .doReturn(sessionToken.v)
            .when(sessionContainer)
            .resolvePartitionLocalSessionToken(Mockito.eq(dsr), Mockito.anyString());

        Mono<StoreResponse> storeResponseSingle =
            consistencyReader.readAsync(dsr, timeoutHelper, false, false);

        StoreResponseValidator validator = StoreResponseValidator
            .create()
            .isSameAs(storeResponse)
            .isSameAs(storeResponse)
            .build();

        validateSuccess(storeResponseSingle, validator);
        transportClientWrapper.verifyNumberOfInvocations(12);
    }

    /**
     * Reading with session consistency, replicas have session token with higher than requested and return not found
     */
    @Test(groups = "unit")
    public void sessionRead_LegitimateNotFound() {
        long lsn = 651175;
        long globalCommittedLsn = 651174;
        String partitionKeyRangeId = "73";

        NotFoundException foundException = new NotFoundException();
        foundException.getResponseHeaders().put(HttpConstants.HttpHeaders.SESSION_TOKEN, partitionKeyRangeId + ":-1#" + lsn);
        foundException.getResponseHeaders().put(WFConstants.BackendHeaders.LSN, Long.toString(lsn));
        foundException.getResponseHeaders().put(WFConstants.BackendHeaders.LOCAL_LSN, Long.toString(lsn));
        foundException.getResponseHeaders().put(WFConstants.BackendHeaders.GLOBAL_COMMITTED_LSN, Long.toString(globalCommittedLsn));

        TransportClientWrapper transportClientWrapper = new TransportClientWrapper.Builder.ReplicaResponseBuilder
                .SequentialBuilder()
                .then(foundException) // 1st replica read returns not found lsn(response) >= lsn(request)
                .then(foundException) // 2nd replica read returns not found lsn(response) >= lsn(request)
                .then(foundException) // 3rd replica read returns not found lsn(response) >= lsn(request)
                .then(foundException) // 4th replica read returns not found lsn(response) >= lsn(request)
                .build();

        Uri primaryUri = Uri.create("primary");
        Uri secondaryUri1 = Uri.create("secondary1");
        Uri secondaryUri2 = Uri.create("secondary2");
        Uri secondaryUri3 = Uri.create("secondary3");

        AddressSelectorWrapper addressSelectorWrapper = AddressSelectorWrapper.Builder.Simple.create()
                .withPrimary(primaryUri)
                .withSecondary(ImmutableList.of(secondaryUri1, secondaryUri2, secondaryUri3))
                .build();
        ISessionContainer sessionContainer = Mockito.mock(ISessionContainer.class);
        TimeoutHelper timeoutHelper = Mockito.mock(TimeoutHelper.class);
        RxDocumentServiceRequest dsr = RxDocumentServiceRequest.createFromName(mockDiagnosticsClientContext(),
                OperationType.Read, "/dbs/db/colls/col/docs/docId", ResourceType.Document);
        dsr.getHeaders().put(HttpConstants.HttpHeaders.CONSISTENCY_LEVEL, ConsistencyLevel.SESSION.toString());
        dsr.requestContext = new DocumentServiceRequestContext();
        Utils.ValueHolder<ISessionToken> sessionToken = Utils.ValueHolder.initialize(null);
        assertThat(VectorSessionToken.tryCreate("-1#" + lsn , sessionToken)).isTrue();
        dsr.requestContext.timeoutHelper = timeoutHelper;
        dsr.requestContext.resolvedPartitionKeyRange = partitionKeyRangeWithId(partitionKeyRangeId);
        dsr.requestContext.requestChargeTracker = new RequestChargeTracker();

        Mockito.doReturn(sessionToken.v).when(sessionContainer).resolvePartitionLocalSessionToken(Mockito.eq(dsr), Mockito.anyString());

        Configs configs = new Configs();
        GatewayServiceConfiguratorReaderMock gatewayServiceConfigurationReaderWrapper = GatewayServiceConfiguratorReaderMock.from(ConsistencyLevel.STRONG,
                                                                                4,
                                                                                3,
                                                                                4,
                                                                                3);

        IAuthorizationTokenProvider authTokenProvider = Mockito.mock(IAuthorizationTokenProvider.class);
        ConsistencyReader consistencyReader = new ConsistencyReader(mockDiagnosticsClientContext(),
                                                                    configs,
                                                                    addressSelectorWrapper.addressSelector,
                                                                    sessionContainer,
                                                                    transportClientWrapper.transportClient,
                                                                    gatewayServiceConfigurationReaderWrapper.gatewayServiceConfigurationReader,
                                                                    authTokenProvider);

        Mono<StoreResponse> storeResponseSingle = consistencyReader.readAsync(dsr, timeoutHelper, false, false);

        FailureValidator failureValidator = FailureValidator.builder().resourceNotFound().instanceOf(NotFoundException.class).unknownSubStatusCode().build();
        validateException(storeResponseSingle, failureValidator);
    }

    /**
     * reading in session consistency, no replica support requested lsn
     */
    @Test(groups = "unit")
    public void sessionRead_ReplicasDoNotHaveTheRequestedLSN() {
        long lsn = 651175;
        long globalCommittedLsn = 651174;
        String partitionKeyRangeId = "73";
        NotFoundException foundException = new NotFoundException();
        foundException.getResponseHeaders().put(HttpConstants.HttpHeaders.SESSION_TOKEN, partitionKeyRangeId + ":-1#" + lsn);
        foundException.getResponseHeaders().put(WFConstants.BackendHeaders.LSN, Long.toString(651175));
        foundException.getResponseHeaders().put(WFConstants.BackendHeaders.LOCAL_LSN, Long.toString(651175));
        foundException.getResponseHeaders().put(WFConstants.BackendHeaders.GLOBAL_COMMITTED_LSN, Long.toString(globalCommittedLsn));

        TransportClientWrapper transportClientWrapper = new TransportClientWrapper.Builder.ReplicaResponseBuilder
                .SequentialBuilder()
                .then(foundException) // 1st replica read lsn lags behind the request lsn
                .then(foundException) // 2nd replica read lsn lags behind the request lsn
                .then(foundException) // 3rd replica read lsn lags behind the request lsn
                .then(foundException) // 4th replica read lsn lags behind the request lsn
                .build();

        Uri primaryUri = Uri.create("primary");
        Uri secondaryUri1 = Uri.create("secondary1");
        Uri secondaryUri2 = Uri.create("secondary2");
        Uri secondaryUri3 = Uri.create("secondary3");

        AddressSelectorWrapper addressSelectorWrapper = AddressSelectorWrapper.Builder.Simple.create()
                .withPrimary(primaryUri)
                .withSecondary(ImmutableList.of(secondaryUri1, secondaryUri2, secondaryUri3))
                .build();
        ISessionContainer sessionContainer = Mockito.mock(ISessionContainer.class);
        TimeoutHelper timeoutHelper = Mockito.mock(TimeoutHelper.class);
        RxDocumentServiceRequest dsr = RxDocumentServiceRequest.createFromName(mockDiagnosticsClientContext(),
                OperationType.Read, "/dbs/db/colls/col/docs/docId", ResourceType.Document);
        dsr.getHeaders().put(HttpConstants.HttpHeaders.CONSISTENCY_LEVEL, ConsistencyLevel.SESSION.toString());
        dsr.requestContext = new DocumentServiceRequestContext();
        Utils.ValueHolder<ISessionToken> sessionToken = Utils.ValueHolder.initialize(null);
        assertThat(VectorSessionToken.tryCreate("-1#" + (lsn + 1) , sessionToken)).isTrue();
        dsr.requestContext.timeoutHelper = timeoutHelper;
        dsr.requestContext.resolvedPartitionKeyRange = partitionKeyRangeWithId(partitionKeyRangeId);
        dsr.requestContext.requestChargeTracker = new RequestChargeTracker();

        Mockito.doReturn(sessionToken.v).when(sessionContainer).resolvePartitionLocalSessionToken(Mockito.eq(dsr), Mockito.anyString());

        Configs configs = new Configs();
        GatewayServiceConfiguratorReaderMock gatewayServiceConfigurationReaderWrapper = GatewayServiceConfiguratorReaderMock.from(ConsistencyLevel.STRONG,
                                                                                4,
                                                                                3,
                                                                                4,
                                                                                3);

        IAuthorizationTokenProvider authTokenProvider = Mockito.mock(IAuthorizationTokenProvider.class);
        ConsistencyReader consistencyReader = new ConsistencyReader(mockDiagnosticsClientContext(),
                                                                    configs,
                                                                    addressSelectorWrapper.addressSelector,
                                                                    sessionContainer,
                                                                    transportClientWrapper.transportClient,
                                                                    gatewayServiceConfigurationReaderWrapper.gatewayServiceConfigurationReader,
                                                                    authTokenProvider);

        Mono<StoreResponse> storeResponseSingle = consistencyReader.readAsync(dsr, timeoutHelper, false, false);

        FailureValidator failureValidator = FailureValidator.builder().resourceNotFound().instanceOf(NotFoundException.class).subStatusCode(HttpConstants.SubStatusCodes.READ_SESSION_NOT_AVAILABLE).build();
        validateException(storeResponseSingle, failureValidator);
    }

    @Test(groups = "unit")
    public void requestRateTooLarge_BubbleUp() {
        long lsn = 651175;
        long globalCommittedLsn = 651174;
        String partitionKeyRangeId = "73";

        RequestRateTooLargeException requestTooLargeException = new RequestRateTooLargeException();
        requestTooLargeException.getResponseHeaders().put(HttpConstants.HttpHeaders.SESSION_TOKEN, partitionKeyRangeId + ":-1#" + lsn);
        requestTooLargeException.getResponseHeaders().put(WFConstants.BackendHeaders.LSN, Long.toString(651175));
        requestTooLargeException.getResponseHeaders().put(WFConstants.BackendHeaders.LOCAL_LSN, Long.toString(651175));
        requestTooLargeException.getResponseHeaders().put(WFConstants.BackendHeaders.GLOBAL_COMMITTED_LSN, Long.toString(globalCommittedLsn));

        TransportClientWrapper transportClientWrapper = new TransportClientWrapper.Builder.ReplicaResponseBuilder
                .SequentialBuilder()
                .then(requestTooLargeException) // 1st replica read result in throttling
                .then(requestTooLargeException) // 2nd replica read result in throttling
                .then(requestTooLargeException) // 3rd replica read result in throttling
                .then(requestTooLargeException) // 4th replica read result in throttling
                .build();

        Uri primaryUri = Uri.create("primary");
        Uri secondaryUri1 = Uri.create("secondary1");
        Uri secondaryUri2 = Uri.create("secondary2");
        Uri secondaryUri3 = Uri.create("secondary3");

        AddressSelectorWrapper addressSelectorWrapper = AddressSelectorWrapper.Builder.Simple.create()
                .withPrimary(primaryUri)
                .withSecondary(ImmutableList.of(secondaryUri1, secondaryUri2, secondaryUri3))
                .build();
        ISessionContainer sessionContainer = Mockito.mock(ISessionContainer.class);
        TimeoutHelper timeoutHelper = Mockito.mock(TimeoutHelper.class);
        RxDocumentServiceRequest dsr = RxDocumentServiceRequest.createFromName(mockDiagnosticsClientContext(),
                OperationType.Read, "/dbs/db/colls/col/docs/docId", ResourceType.Document);
        dsr.getHeaders().put(HttpConstants.HttpHeaders.CONSISTENCY_LEVEL, ConsistencyLevel.SESSION.toString());
        dsr.requestContext = new DocumentServiceRequestContext();
        Utils.ValueHolder<ISessionToken> sessionToken = Utils.ValueHolder.initialize(null);
        assertThat(VectorSessionToken.tryCreate("-1#" + lsn , sessionToken)).isTrue();
        dsr.requestContext.timeoutHelper = timeoutHelper;
        dsr.requestContext.resolvedPartitionKeyRange = partitionKeyRangeWithId(partitionKeyRangeId);
        dsr.requestContext.requestChargeTracker = new RequestChargeTracker();

        Mockito.doReturn(sessionToken.v).when(sessionContainer).resolvePartitionLocalSessionToken(Mockito.eq(dsr), Mockito.anyString());

        Configs configs = new Configs();
        GatewayServiceConfiguratorReaderMock gatewayServiceConfigurationReaderWrapper = GatewayServiceConfiguratorReaderMock.from(ConsistencyLevel.STRONG,
                                                                                4,
                                                                                3,
                                                                                4,
                                                                                3);

        IAuthorizationTokenProvider authTokenProvider = Mockito.mock(IAuthorizationTokenProvider.class);
        ConsistencyReader consistencyReader = new ConsistencyReader(mockDiagnosticsClientContext(),
                                                                    configs,
                                                                    addressSelectorWrapper.addressSelector,
                                                                    sessionContainer,
                                                                    transportClientWrapper.transportClient,
                                                                    gatewayServiceConfigurationReaderWrapper.gatewayServiceConfigurationReader,
                                                                    authTokenProvider);

        Mono<StoreResponse> storeResponseSingle = consistencyReader.readAsync(dsr, timeoutHelper, false, false);

        FailureValidator failureValidator = FailureValidator.builder().instanceOf(RequestRateTooLargeException.class).unknownSubStatusCode().build();
        validateException(storeResponseSingle, failureValidator);
    }

    @DataProvider(name = "simpleReadStrongArgProvider")
    public Object[][] simpleReadStrongArgProvider() {
        return new Object[][]{
                {  1, ReadMode.Strong },
                {  2, ReadMode.Strong },
                {  3, ReadMode.Strong },
        };
    }

    @Test(groups = "unit", dataProvider = "simpleReadStrongArgProvider")
    public void basicReadStrong_AllReplicasSameLSN(int replicaCountToRead, ReadMode readMode) {
        ISessionContainer sessionContainer = Mockito.mock(ISessionContainer.class);
        Uri primaryReplicaURI = Uri.create("primary");
        ImmutableList<Uri> secondaryReplicaURIs = ImmutableList.of(Uri.create("secondary1"), Uri.create("secondary2"), Uri.create("secondary3"));
        AddressSelectorWrapper addressSelectorWrapper = AddressSelectorWrapper.Builder.Simple.create()
                .withPrimary(primaryReplicaURI)
                .withSecondary(secondaryReplicaURIs)
                .build();

        RxDocumentServiceRequest request = RxDocumentServiceRequest.createFromName(mockDiagnosticsClientContext(),
                OperationType.Read, "/dbs/db/colls/col/docs/docId", ResourceType.Document);

        request.requestContext = new DocumentServiceRequestContext();
        request.requestContext.timeoutHelper = Mockito.mock(TimeoutHelper.class);
        request.requestContext.resolvedPartitionKeyRange = Mockito.mock(PartitionKeyRange.class);
        request.requestContext.requestChargeTracker = new RequestChargeTracker();

        BigDecimal requestChargePerRead = new BigDecimal(1.1);

        StoreResponse primaryResponse = StoreResponseBuilder.create()
                .withLSN(51)
                .withLocalLSN(18)
                .withRequestCharge(requestChargePerRead.doubleValue())
                .build();
        StoreResponse secondaryResponse1 = StoreResponseBuilder.create()
                .withLSN(51)
                .withLocalLSN(18)
                .withRequestCharge(requestChargePerRead.doubleValue())
                .build();
        StoreResponse secondaryResponse2 = StoreResponseBuilder.create()
                .withLSN(51)
                .withLocalLSN(18)
                .withRequestCharge(requestChargePerRead.doubleValue())
                .build();
        StoreResponse secondaryResponse3 = StoreResponseBuilder.create()
                .withLSN(51)
                .withLocalLSN(18)
                .withRequestCharge(requestChargePerRead.doubleValue())
                .build();

        TransportClientWrapper transportClientWrapper = TransportClientWrapper.Builder.uriToResultBuilder()
                .storeResponseOn(primaryReplicaURI, OperationType.Read, ResourceType.Document, primaryResponse, false)
                .storeResponseOn(secondaryReplicaURIs.get(0), OperationType.Read, ResourceType.Document, secondaryResponse1, false)
                .storeResponseOn(secondaryReplicaURIs.get(1), OperationType.Read, ResourceType.Document, secondaryResponse2, false)
                .storeResponseOn(secondaryReplicaURIs.get(2), OperationType.Read, ResourceType.Document, secondaryResponse3, false)
                .build();

        StoreReader storeReader = new StoreReader(transportClientWrapper.transportClient, addressSelectorWrapper.addressSelector, sessionContainer);
        GatewayServiceConfigurationReader serviceConfigurator = Mockito.mock(GatewayServiceConfigurationReader.class);
        IAuthorizationTokenProvider authTokenProvider = Mockito.mock(IAuthorizationTokenProvider.class);
        QuorumReader quorumReader = new QuorumReader(mockDiagnosticsClientContext(), configs, transportClientWrapper.transportClient, addressSelectorWrapper.addressSelector, storeReader, serviceConfigurator, authTokenProvider);

        Mono<StoreResponse> storeResponseSingle = quorumReader.readStrongAsync(mockDiagnosticsClientContext(), request, replicaCountToRead, readMode);

        StoreResponseValidator validator = StoreResponseValidator.create()
                .withBELSN(51)
                .withRequestCharge(requestChargePerRead.multiply(BigDecimal.valueOf(replicaCountToRead)).setScale(2, RoundingMode.FLOOR).doubleValue())
                .build();
        validateSuccess(storeResponseSingle, validator);

        transportClientWrapper.validate()
                .verifyNumberOfInvocations(replicaCountToRead);
        addressSelectorWrapper.validate()
                .verifyNumberOfForceCachRefresh(0)
                .verifyVesolvePrimaryUriAsyncCount(0)
                .verifyTotalInvocations(1);
    }

    // TODO: add more mocking tests for when one replica lags behind and we need to do barrier request.

    public static void validateSuccess(Mono<List<StoreResult>> single,
                                       MultiStoreResultValidator validator) {
        validateSuccess(single, validator, 10000);
    }

    public static void validateSuccess(Mono<List<StoreResult>> single,
                                       MultiStoreResultValidator validator,
                                       long timeout) {
        TestSubscriber<List<StoreResult>> testSubscriber = new TestSubscriber<>();

        single.subscribe(testSubscriber);
        testSubscriber.awaitTerminalEvent(timeout, TimeUnit.MILLISECONDS);
        testSubscriber.assertNoErrors();
        testSubscriber.assertComplete();
        testSubscriber.assertValueCount(1);
        validator.validate(testSubscriber.values().get(0));
    }

    public static void validateSuccess(Mono<StoreResponse> single,
                                       StoreResponseValidator validator) {
        validateSuccess(single, validator, 10000);
    }

    public static void validateSuccess(Mono<StoreResponse> single,
                                       StoreResponseValidator validator,
                                       long timeout) {
        TestSubscriber<StoreResponse> testSubscriber = new TestSubscriber<>();

        single.subscribe(testSubscriber);
        testSubscriber.awaitTerminalEvent(timeout, TimeUnit.MILLISECONDS);
        testSubscriber.assertNoErrors();
        testSubscriber.assertComplete();
        testSubscriber.assertValueCount(1);
        validator.validate(testSubscriber.values().get(0));
    }

    public static <T> void validateException(Mono<T> single,
                                             FailureValidator validator,
                                             long timeout) {
        TestSubscriber<T> testSubscriber = new TestSubscriber<>();

        single.subscribe(testSubscriber);
        testSubscriber.awaitTerminalEvent(timeout, TimeUnit.MILLISECONDS);
        testSubscriber.assertNotComplete();
        testSubscriber.assertTerminated();
        assertThat(testSubscriber.errorCount()).isEqualTo(1);
        validator.validate(testSubscriber.errors().get(0));
    }

    public static <T> void validateException(Mono<T> single,
                                             FailureValidator validator) {
        validateException(single, validator, TIMEOUT);
    }

    private PartitionKeyRange partitionKeyRangeWithId(String id) {
        PartitionKeyRange partitionKeyRange = Mockito.mock(PartitionKeyRange.class);
        Mockito.doReturn(id).when(partitionKeyRange).getId();
        return partitionKeyRange;
    }
}


text:
import re
import testClasses
import textwrap

# import project specific code
import layout
import pacman
from search import SearchProblem

# helper function for printing solutions in solution files
def wrap_solution(solution):
    if type(solution) == type([]):
        return '\n'.join(textwrap.wrap(' '.join(solution)))
    else:
        return str(solution)

def followAction(state, action, problem):
  for successor1, action1, cost1 in problem.getSuccessors(state):
    if action == action1: return successor1
  return None

def followPath(path, problem):
  state = problem.getStartState()
  states = [state]
  for action in path:
    state = followAction(state, action, problem)
    states.append(state)
  return states

def checkSolution(problem, path):
  state = problem.getStartState()
  for action in path:
    state = followAction(state, action, problem)
  return problem.isGoalState(state)

# Search problem on a plain graph
class GraphSearch(SearchProblem):

    # Read in the state graph; define start/end states, edges and costs
    def __init__(self, graph_text):
        self.expanded_states = []
        lines = graph_text.split('\n')
        r = re.match('start_state:(.*)', lines[0])
        if r == None:
            print "Broken graph:"
            print '"""%s"""' % graph_text
            raise Exception("GraphSearch graph specification start_state not found or incorrect on line:" + l)
        self.start_state = r.group(1).strip()
        r = re.match('goal_states:(.*)', lines[1])
        if r == None:
            print "Broken graph:"
            print '"""%s"""' % graph_text
            raise Exception("GraphSearch graph specification goal_states not found or incorrect on line:" + l)
        goals = r.group(1).split()
        self.goals = map(str.strip, goals)
        self.successors = {}
        all_states = set()
        self.orderedSuccessorTuples = []
        for l in lines[2:]:
            if len(l.split()) == 3:
                start, action, next_state = l.split()
                cost = 1
            elif len(l.split()) == 4:
                start, action, next_state, cost = l.split()
            else:
                print "Broken graph:"
                print '"""%s"""' % graph_text
                raise Exception("Invalid line in GraphSearch graph specification on line:" + l)
            cost = float(cost)
            self.orderedSuccessorTuples.append((start, action, next_state, cost))
            all_states.add(start)
            all_states.add(next_state)
            if start not in self.successors:
                self.successors[start] = []
            self.successors[start].append((next_state, action, cost))
        for s in all_states:
            if s not in self.successors:
                self.successors[s] = []

    # Get start state
    def getStartState(self):
        return self.start_state

    # Check if a state is a goal state
    def isGoalState(self, state):
        return state in self.goals

    # Get all successors of a state
    def getSuccessors(self, state):
        self.expanded_states.append(state)
        return list(self.successors[state])

    # Calculate total cost of a sequence of actions
    def getCostOfActions(self, actions):
        total_cost = 0
        state = self.start_state
        for a in actions:
            successors = self.successors[state]
            match = False
            for (next_state, action, cost) in successors:
                if a == action:
                    state = next_state
                    total_cost += cost
                    match = True
            if not match:
                print 'invalid action sequence'
                sys.exit(1)
        return total_cost

    # Return a list of all states on which 'getSuccessors' was called
    def getExpandedStates(self):
        return self.expanded_states

    def __str__(self):
        print self.successors
        edges = ["%s %s %s %s" % t for t in self.orderedSuccessorTuples]
        return \
"""start_state: %s
goal_states: %s
%s""" % (self.start_state, " ".join(self.goals), "\n".join(edges))

def parseHeuristic(heuristicText):
    heuristic = {}
    for line in heuristicText.split('\n'):
        tokens = line.split()
        if len(tokens) != 2:
            print "Broken heuristic:"
            print '"""%s"""' % graph_text
            raise Exception("GraphSearch heuristic specification broken:" + l)
        state, h = tokens
        heuristic[state] = float(h)

    def graphHeuristic(state, problem=None):
        if state in heuristic:
            return heuristic[state]
        else:
            pp = pprint.PrettyPrinter(indent=4)
            print "Heuristic:"
            pp.pprint(heuristic)
            raise Exception("Graph heuristic called with invalid state: " + str(state))

    return graphHeuristic

class GraphSearchTest(testClasses.TestCase):

    def __init__(self, question, testDict):
        super(GraphSearchTest, self).__init__(question, testDict)
        self.graph_text = testDict['graph']
        self.alg = testDict['algorithm']
        self.diagram = testDict['diagram']
        self.exactExpansionOrder = testDict.get('exactExpansionOrder', 'True').lower() == "true"
        if 'heuristic' in testDict:
            self.heuristic = parseHeuristic(testDict['heuristic'])
        else:
            self.heuristic = None

    # Note that the return type of this function is a tripple:
    # (solution, expanded states, error message)
    def getSolInfo(self, search):
        alg = getattr(search, self.alg)
        problem = GraphSearch(self.graph_text)
        if self.heuristic != None:
            solution = alg(problem, self.heuristic)
        else:
            solution = alg(problem)

        if type(solution) != type([]):
            return None, None, 'The result of %s must be a list. (Instead, it is %s)' % (self.alg, type(solution))

        return solution, problem.getExpandedStates(), None

    # Run student code.  If an error message is returned, print error and return false.
    # If a good solution is returned, printn the solution and return true; otherwise,
    # print both the correct and student's solution and return false.
    def execute(self, grades, moduleDict, solutionDict):
        search = moduleDict['search']
        searchAgents = moduleDict['searchAgents']
        gold_solution = [str.split(solutionDict['solution']), str.split(solutionDict['rev_solution'])]
        gold_expanded_states = [str.split(solutionDict['expanded_states']), str.split(solutionDict['rev_expanded_states'])]

        solution, expanded_states, error = self.getSolInfo(search)
        if error != None:
            grades.addMessage('FAIL: %s' % self.path)
            grades.addMessage('\t%s' % error)
            return False

        if solution in gold_solution and (not self.exactExpansionOrder or expanded_states in gold_expanded_states):
            grades.addMessage('PASS: %s' % self.path)
            grades.addMessage('\tsolution:\t\t%s' % solution)
            grades.addMessage('\texpanded_states:\t%s' % expanded_states)
            return True
        else:
            grades.addMessage('FAIL: %s' % self.path)
            grades.addMessage('\tgraph:')
            for line in self.diagram.split('\n'):
                grades.addMessage('\t    %s' % (line,))
            grades.addMessage('\tstudent solution:\t\t%s' % solution)
            grades.addMessage('\tstudent expanded_states:\t%s' % expanded_states)
            grades.addMessage('')
            grades.addMessage('\tcorrect solution:\t\t%s' % gold_solution[0])
            grades.addMessage('\tcorrect expanded_states:\t%s' % gold_expanded_states[0])
            grades.addMessage('\tcorrect rev_solution:\t\t%s' % gold_solution[1])
            grades.addMessage('\tcorrect rev_expanded_states:\t%s' % gold_expanded_states[1])
            return False

    def writeSolution(self, moduleDict, filePath):
        search = moduleDict['search']
        searchAgents = moduleDict['searchAgents']
        # open file and write comments
        handle = open(filePath, 'w')
        handle.write('# This is the solution file for %s.\n' % self.path)
        handle.write('# This solution is designed to support both right-to-left\n')
        handle.write('# and left-to-right implementations.\n')

        # write forward solution
        solution, expanded_states, error = self.getSolInfo(search)
        if error != None: raise Exception("Error in solution code: %s" % error)
        handle.write('solution: "%s"\n' % ' '.join(solution))
        handle.write('expanded_states: "%s"\n' % ' '.join(expanded_states))

        # reverse and write backwards solution
        search.REVERSE_PUSH = not search.REVERSE_PUSH
        solution, expanded_states, error = self.getSolInfo(search)
        if error != None: raise Exception("Error in solution code: %s" % error)
        handle.write('rev_solution: "%s"\n' % ' '.join(solution))
        handle.write('rev_expanded_states: "%s"\n' % ' '.join(expanded_states))

        # clean up
        search.REVERSE_PUSH = not search.REVERSE_PUSH
        handle.close()
        return True

class PacmanSearchTest(testClasses.TestCase):

    def __init__(self, question, testDict):
        super(PacmanSearchTest, self).__init__(question, testDict)
        self.layout_text = testDict['layout']
        self.alg = testDict['algorithm']
        self.layoutName = testDict['layoutName']

        # TODO: sensible to have defaults like this?
        self.leewayFactor = float(testDict.get('leewayFactor', '1'))
        self.costFn = eval(testDict.get('costFn', 'None'))
        self.searchProblemClassName = testDict.get('searchProblemClass', 'PositionSearchProblem')
        self.heuristicName = testDict.get('heuristic', None)

    def getSolInfo(self, search, searchAgents):
        alg = getattr(search, self.alg)
        lay = layout.Layout([l.strip() for l in self.layout_text.split('\n')])
        start_state = pacman.GameState()
        start_state.initialize(lay, 0)

        problemClass = getattr(searchAgents, self.searchProblemClassName)
        problemOptions = {}
        if self.costFn != None:
            problemOptions['costFn'] = self.costFn
        problem = problemClass(start_state, **problemOptions)
        heuristic = getattr(searchAgents, self.heuristicName) if self.heuristicName != None else None

        if heuristic != None:
            solution = alg(problem, heuristic)
        else:
            solution = alg(problem)

        if type(solution) != type([]):
            return None, None, 'The result of %s must be a list. (Instead, it is %s)' % (self.alg, type(solution))

        from game import Directions
        dirs = Directions.LEFT.keys()
        if [el in dirs for el in solution].count(False) != 0:
            return None, None, 'Output of %s must be a list of actions from game.Directions' % self.alg

        expanded = problem._expanded
        return solution, expanded, None

    def execute(self, grades, moduleDict, solutionDict):
        search = moduleDict['search']
        searchAgents = moduleDict['searchAgents']
        gold_solution = [str.split(solutionDict['solution']), str.split(solutionDict['rev_solution'])]
        gold_expanded = max(int(solutionDict['expanded_nodes']), int(solutionDict['rev_expanded_nodes']))

        solution, expanded, error = self.getSolInfo(search, searchAgents)
        if error != None:
            grades.addMessage('FAIL: %s' % self.path)
            grades.addMessage('%s' % error)
            return False

        # FIXME: do we want to standardize test output format?

        if solution not in gold_solution:
            grades.addMessage('FAIL: %s' % self.path)
            grades.addMessage('Solution not correct.')
            grades.addMessage('\tstudent solution length: %s' % len(solution))
            grades.addMessage('\tstudent solution:\n%s' % wrap_solution(solution))
            grades.addMessage('')
            grades.addMessage('\tcorrect solution length: %s' % len(gold_solution[0]))
            grades.addMessage('\tcorrect (reversed) solution length: %s' % len(gold_solution[1]))
            grades.addMessage('\tcorrect solution:\n%s' % wrap_solution(gold_solution[0]))
            grades.addMessage('\tcorrect (reversed) solution:\n%s' % wrap_solution(gold_solution[1]))
            return False

        if expanded > self.leewayFactor * gold_expanded and expanded > gold_expanded + 1:
            grades.addMessage('FAIL: %s' % self.path)
            grades.addMessage('Too many node expanded; are you expanding nodes twice?')
            grades.addMessage('\tstudent nodes expanded: %s' % expanded)
            grades.addMessage('')
            grades.addMessage('\tcorrect nodes expanded: %s (leewayFactor %s)' % (gold_expanded, self.leewayFactor))
            return False

        grades.addMessage('PASS: %s' % self.path)
        grades.addMessage('\tpacman layout:\t\t%s' % self.layoutName)
        grades.addMessage('\tsolution length: %s' % len(solution))
        grades.addMessage('\tnodes expanded:\t\t%s' % expanded)
        return True

    def writeSolution(self, moduleDict, filePath):
        search = moduleDict['search']
        searchAgents = moduleDict['searchAgents']
        # open file and write comments
        handle = open(filePath, 'w')
        handle.write('# This is the solution file for %s.\n' % self.path)
        handle.write('# This solution is designed to support both right-to-left\n')
        handle.write('# and left-to-right implementations.\n')
        handle.write('# Number of nodes expanded must be with a factor of %s of the numbers below.\n' % self.leewayFactor)

        # write forward solution
        solution, expanded, error = self.getSolInfo(search, searchAgents)
        if error != None: raise Exception("Error in solution code: %s" % error)
        handle.write('solution: """\n%s\n"""\n' % wrap_solution(solution))
        handle.write('expanded_nodes: "%s"\n' % expanded)

        # write backward solution
        search.REVERSE_PUSH = not search.REVERSE_PUSH
        solution, expanded, error = self.getSolInfo(search, searchAgents)
        if error != None: raise Exception("Error in solution code: %s" % error)
        handle.write('rev_solution: """\n%s\n"""\n' % wrap_solution(solution))
        handle.write('rev_expanded_nodes: "%s"\n' % expanded)

        # clean up
        search.REVERSE_PUSH = not search.REVERSE_PUSH
        handle.close()
        return True

from game import Actions
def getStatesFromPath(start, path):
    "Returns the list of states visited along the path"
    vis = [start]
    curr = start
    for a in path:
        x,y = curr
        dx, dy = Actions.directionToVector(a)
        curr = (int(x + dx), int(y + dy))
        vis.append(curr)
    return vis

class CornerProblemTest(testClasses.TestCase):

    def __init__(self, question, testDict):
        super(CornerProblemTest, self).__init__(question, testDict)
        self.layoutText = testDict['layout']
        self.layoutName = testDict['layoutName']

    def solution(self, search, searchAgents):
        lay = layout.Layout([l.strip() for l in self.layoutText.split('\n')])
        gameState = pacman.GameState()
        gameState.initialize(lay, 0)
        problem = searchAgents.CornersProblem(gameState)
        path = search.bfs(problem)

        gameState = pacman.GameState()
        gameState.initialize(lay, 0)
        visited = getStatesFromPath(gameState.getPacmanPosition(), path)
        top, right = gameState.getWalls().height-2, gameState.getWalls().width-2
        missedCorners = [p for p in ((1,1), (1,top), (right, 1), (right, top)) if p not in visited]

        return path, missedCorners

    def execute(self, grades, moduleDict, solutionDict):
        search = moduleDict['search']
        searchAgents = moduleDict['searchAgents']
        gold_length = int(solutionDict['solution_length'])
        solution, missedCorners = self.solution(search, searchAgents)

        if type(solution) != type([]):
            grades.addMessage('FAIL: %s' % self.path)
            grades.addMessage('The result must be a list. (Instead, it is %s)' % type(solution))
            return False

        if len(missedCorners) != 0:
            grades.addMessage('FAIL: %s' % self.path)
            grades.addMessage('Corners missed: %s' % missedCorners)
            return False

        if len(solution) != gold_length:
            grades.addMessage('FAIL: %s' % self.path)
            grades.addMessage('Optimal solution not found.')
            grades.addMessage('\tstudent solution length:\n%s' % len(solution))
            grades.addMessage('')
            grades.addMessage('\tcorrect solution length:\n%s' % gold_length)
            return False

        grades.addMessage('PASS: %s' % self.path)
        grades.addMessage('\tpacman layout:\t\t%s' % self.layoutName)
        grades.addMessage('\tsolution length:\t\t%s' % len(solution))
        return True

    def writeSolution(self, moduleDict, filePath):
        search = moduleDict['search']
        searchAgents = moduleDict['searchAgents']
        # open file and write comments
        handle = open(filePath, 'w')
        handle.write('# This is the solution file for %s.\n' % self.path)

        print "Solving problem", self.layoutName
        print self.layoutText

        path, _ = self.solution(search, searchAgents)
        length = len(path)
        print "Problem solved"

        handle.write('solution_length: "%s"\n' % length)
        handle.close()

# template = """class: "HeuristicTest"
#
# heuristic: "foodHeuristic"
# searchProblemClass: "FoodSearchProblem"
# layoutName: "Test %s"
# layout: \"\"\"
# %s
# \"\"\"
# """
#
# for i, (_, _, l) in enumerate(doneTests + foodTests):
#     f = open("food_heuristic_%s.test" % (i+1), "w")
#     f.write(template % (i+1, "\n".join(l)))
#     f.close()

class HeuristicTest(testClasses.TestCase):

    def __init__(self, question, testDict):
        super(HeuristicTest, self).__init__(question, testDict)
        self.layoutText = testDict['layout']
        self.layoutName = testDict['layoutName']
        self.searchProblemClassName = testDict['searchProblemClass']
        self.heuristicName = testDict['heuristic']

    def setupProblem(self, searchAgents):
        lay = layout.Layout([l.strip() for l in self.layoutText.split('\n')])
        gameState = pacman.GameState()
        gameState.initialize(lay, 0)
        problemClass = getattr(searchAgents, self.searchProblemClassName)
        problem = problemClass(gameState)
        state = problem.getStartState()
        heuristic = getattr(searchAgents, self.heuristicName)

        return problem, state, heuristic

    def checkHeuristic(self, heuristic, problem, state, solutionCost):
        h0 = heuristic(state, problem)

        if solutionCost == 0:
            if h0 == 0:
                return True, ''
            else:
                return False, 'Heuristic failed H(goal) == 0 test'

        if h0 < 0:
            return False, 'Heuristic failed H >= 0 test'
        if not h0 > 0:
            return False, 'Heuristic failed non-triviality test'
        if not h0 <= solutionCost:
            return False, 'Heuristic failed admissibility test'

        for succ, action, stepCost in problem.getSuccessors(state):
            h1 = heuristic(succ, problem)
            if h1 < 0: return False, 'Heuristic failed H >= 0 test'
            if h0 - h1 > stepCost: return False, 'Heuristic failed consistency test'

        return True, ''

    def execute(self, grades, moduleDict, solutionDict):
        search = moduleDict['search']
        searchAgents = moduleDict['searchAgents']
        solutionCost = int(solutionDict['solution_cost'])
        problem, state, heuristic = self.setupProblem(searchAgents)

        passed, message = self.checkHeuristic(heuristic, problem, state, solutionCost)

        if not passed:
            grades.addMessage('FAIL: %s' % self.path)
            grades.addMessage('%s' % message)
            return False
        else:
            grades.addMessage('PASS: %s' % self.path)
            return True

    def writeSolution(self, moduleDict, filePath):
        search = moduleDict['search']
        searchAgents = moduleDict['searchAgents']
        # open file and write comments
        handle = open(filePath, 'w')
        handle.write('# This is the solution file for %s.\n' % self.path)

        print "Solving problem", self.layoutName, self.heuristicName
        print self.layoutText
        problem, _, heuristic = self.setupProblem(searchAgents)
        path = search.astar(problem, heuristic)
        cost = problem.getCostOfActions(path)
        print "Problem solved"

        handle.write('solution_cost: "%s"\n' % cost)
        handle.close()
        return True

class HeuristicGrade(testClasses.TestCase):

    def __init__(self, question, testDict):
        super(HeuristicGrade, self).__init__(question, testDict)
        self.layoutText = testDict['layout']
        self.layoutName = testDict['layoutName']
        self.searchProblemClassName = testDict['searchProblemClass']
        self.heuristicName = testDict['heuristic']
        self.basePoints = int(testDict['basePoints'])
        self.thresholds = [int(t) for t in testDict['gradingThresholds'].split()]

    def setupProblem(self, searchAgents):
        lay = layout.Layout([l.strip() for l in self.layoutText.split('\n')])
        gameState = pacman.GameState()
        gameState.initialize(lay, 0)
        problemClass = getattr(searchAgents, self.searchProblemClassName)
        problem = problemClass(gameState)
        state = problem.getStartState()
        heuristic = getattr(searchAgents, self.heuristicName)

        return problem, state, heuristic

    def execute(self, grades, moduleDict, solutionDict):
        search = moduleDict['search']
        searchAgents = moduleDict['searchAgents']
        problem, _, heuristic = self.setupProblem(searchAgents)

        path = search.astar(problem, heuristic)

        expanded = problem._expanded

        if not checkSolution(problem, path):
            grades.addMessage('FAIL: %s' % self.path)
            grades.addMessage('\tReturned path is not a solution.')
            grades.addMessage('\tpath returned by astar: %s' % expanded)
            return False

        grades.addPoints(self.basePoints)
        points = 0
        for threshold in self.thresholds:
            if expanded <= threshold:
                points += 1
        grades.addPoints(points)
        if points >= len(self.thresholds):
            grades.addMessage('PASS: %s' % self.path)
        else:
            grades.addMessage('FAIL: %s' % self.path)
        grades.addMessage('\texpanded nodes: %s' % expanded)
        grades.addMessage('\tthresholds: %s' % self.thresholds)

        return True

    def writeSolution(self, moduleDict, filePath):
        handle = open(filePath, 'w')
        handle.write('# This is the solution file for %s.\n' % self.path)
        handle.write('# File intentionally blank.\n')
        handle.close()
        return True

# template = """class: "ClosestDotTest"
#
# layoutName: "Test %s"
# layout: \"\"\"
# %s
# \"\"\"
# """
#
# for i, (_, _, l) in enumerate(foodTests):
#     f = open("closest_dot_%s.test" % (i+1), "w")
#     f.write(template % (i+1, "\n".join(l)))
#     f.close()

class ClosestDotTest(testClasses.TestCase):

    def __init__(self, question, testDict):
        super(ClosestDotTest, self).__init__(question, testDict)
        self.layoutText = testDict['layout']
        self.layoutName = testDict['layoutName']

    def solution(self, searchAgents):
        lay = layout.Layout([l.strip() for l in self.layoutText.split('\n')])
        gameState = pacman.GameState()
        gameState.initialize(lay, 0)
        path = searchAgents.ClosestDotSearchAgent().findPathToClosestDot(gameState)
        return path

    def execute(self, grades, moduleDict, solutionDict):
        search = moduleDict['search']
        searchAgents = moduleDict['searchAgents']
        gold_length = int(solutionDict['solution_length'])
        solution = self.solution(searchAgents)

        if type(solution) != type([]):
            grades.addMessage('FAIL: %s' % self.path)
            grades.addMessage('\tThe result must be a list. (Instead, it is %s)' % type(solution))
            return False

        if len(solution) != gold_length:
            grades.addMessage('FAIL: %s' % self.path)
            grades.addMessage('Closest dot not found.')
            grades.addMessage('\tstudent solution length:\n%s' % len(solution))
            grades.addMessage('')
            grades.addMessage('\tcorrect solution length:\n%s' % gold_length)
            return False

        grades.addMessage('PASS: %s' % self.path)
        grades.addMessage('\tpacman layout:\t\t%s' % self.layoutName)
        grades.addMessage('\tsolution length:\t\t%s' % len(solution))
        return True

    def writeSolution(self, moduleDict, filePath):
        search = moduleDict['search']
        searchAgents = moduleDict['searchAgents']
        # open file and write comments
        handle = open(filePath, 'w')
        handle.write('# This is the solution file for %s.\n' % self.path)

        print "Solving problem", self.layoutName
        print self.layoutText

        length = len(self.solution(searchAgents))
        print "Problem solved"

        handle.write('solution_length: "%s"\n' % length)
        handle.close()
        return True

class CornerHeuristicSanity(testClasses.TestCase):

    def __init__(self, question, testDict):
        super(CornerHeuristicSanity, self).__init__(question, testDict)
        self.layout_text = testDict['layout']

    def execute(self, grades, moduleDict, solutionDict):
        search = moduleDict['search']
        searchAgents = moduleDict['searchAgents']
        game_state = pacman.GameState()
        lay = layout.Layout([l.strip() for l in self.layout_text.split('\n')])
        game_state.initialize(lay, 0)
        problem = searchAgents.CornersProblem(game_state)
        start_state = problem.getStartState()
        h0 = searchAgents.cornersHeuristic(start_state, problem)
        succs = problem.getSuccessors(start_state)
        # cornerConsistencyA
        for succ in succs:
            h1 = searchAgents.cornersHeuristic(succ[0], problem)
            if h0 - h1 > 1:
                grades.addMessage('FAIL: inconsistent heuristic')
                return False
        heuristic_cost = searchAgents.cornersHeuristic(start_state, problem)
        true_cost = float(solutionDict['cost'])
        # cornerNontrivial
        if heuristic_cost == 0:
            grades.addMessage('FAIL: must use non-trivial heuristic')
            return False
        # cornerAdmissible
        if heuristic_cost > true_cost:
            grades.addMessage('FAIL: Inadmissible heuristic')
            return False
        path = solutionDict['path'].split()
        states = followPath(path, problem)
        heuristics = []
        for state in states:
            heuristics.append(searchAgents.cornersHeuristic(state, problem))
        for i in range(0, len(heuristics) - 1):
            h0 = heuristics[i]
            h1 = heuristics[i+1]
            # cornerConsistencyB
            if h0 - h1 > 1:
                grades.addMessage('FAIL: inconsistent heuristic')
                return False
            # cornerPosH
            if h0 < 0 or h1 <0:
                grades.addMessage('FAIL: non-positive heuristic')
                return False
        # cornerGoalH
        if heuristics[len(heuristics) - 1] != 0:
            grades.addMessage('FAIL: heuristic non-zero at goal')
            return False
        grades.addMessage('PASS: heuristic value less than true cost at start state')
        return True

    def writeSolution(self, moduleDict, filePath):
        search = moduleDict['search']
        searchAgents = moduleDict['searchAgents']
        # write comment
        handle = open(filePath, 'w')
        handle.write('# In order for a heuristic to be admissible, the value\n')
        handle.write('# of the heuristic must be less at each state than the\n')
        handle.write('# true cost of the optimal path from that state to a goal.\n')

        # solve problem and write solution
        lay = layout.Layout([l.strip() for l in self.layout_text.split('\n')])
        start_state = pacman.GameState()
        start_state.initialize(lay, 0)
        problem = searchAgents.CornersProblem(start_state)
        solution = search.astar(problem, searchAgents.cornersHeuristic)
        handle.write('cost: "%d"\n' % len(solution))
        handle.write('path: """\n%s\n"""\n' % wrap_solution(solution))
        handle.close()
        return True

class CornerHeuristicPacman(testClasses.TestCase):

    def __init__(self, question, testDict):
        super(CornerHeuristicPacman, self).__init__(question, testDict)
        self.layout_text = testDict['layout']

    def execute(self, grades, moduleDict, solutionDict):
        search = moduleDict['search']
        searchAgents = moduleDict['searchAgents']
        total = 0
        true_cost = float(solutionDict['cost'])
        thresholds = map(int, solutionDict['thresholds'].split())
        game_state = pacman.GameState()
        lay = layout.Layout([l.strip() for l in self.layout_text.split('\n')])
        game_state.initialize(lay, 0)
        problem = searchAgents.CornersProblem(game_state)
        start_state = problem.getStartState()
        if searchAgents.cornersHeuristic(start_state, problem) > true_cost:
            grades.addMessage('FAIL: Inadmissible heuristic')
            return False
        path = search.astar(problem, searchAgents.cornersHeuristic)
        print "path:", path
        print "path length:", len(path)
        cost = problem.getCostOfActions(path)
        if cost > true_cost:
            grades.addMessage('FAIL: Inconsistent heuristic')
            return False
        expanded = problem._expanded
        points = 0
        for threshold in thresholds:
            if expanded <= threshold:
                points += 1
        grades.addPoints(points)
        if points >= len(thresholds):
            grades.addMessage('PASS: Heuristic resulted in expansion of %d nodes' % expanded)
        else:
            grades.addMessage('FAIL: Heuristic resulted in expansion of %d nodes' % expanded)
        return True

    def writeSolution(self, moduleDict, filePath):
        search = moduleDict['search']
        searchAgents = moduleDict['searchAgents']
        # write comment
        handle = open(filePath, 'w')
        handle.write('# This solution file specifies the length of the optimal path\n')
        handle.write('# as well as the thresholds on number of nodes expanded to be\n')
        handle.write('# used in scoring.\n')

        # solve problem and write solution
        lay = layout.Layout([l.strip() for l in self.layout_text.split('\n')])
        start_state = pacman.GameState()
        start_state.initialize(lay, 0)
        problem = searchAgents.CornersProblem(start_state)
        solution = search.astar(problem, searchAgents.cornersHeuristic)
        handle.write('cost: "%d"\n' % len(solution))
        handle.write('path: """\n%s\n"""\n' % wrap_solution(solution))
        handle.write('thresholds: "2000 1600 1200"\n')
        handle.close()
        return True

redpajama_set_name: RedPajamaGithub


text:
package leap.web.action;

import leap.lang.reflect.ReflectMethod;

import java.lang.annotation.Annotation;
import java.lang.reflect.Type;
import java.util.Map;

public class ActionWrapper implements Action {

    protected final Action action;

    public ActionWrapper(Action action) {
        this.action = action;
    }

    @Override
    public Object execute(ActionContext context, Object[] args) throws Throwable {
        return action.execute(context, args);
    }

    @Override
    public <T> void setExtension(Class<T> type, Object extension) {
        action.setExtension(type, extension);
    }

    @Override
    public <T> T removeExtension(Class<?> type) {
        return action.removeExtension(type);
    }

    @Override
    public Map<Class<?>, Object> getExtensions() {
        return action.getExtensions();
    }

    @Override
    public <T> T getExtension(Class<?> type) {
        return action.getExtension(type);
    }

    @Override
    public String getName() {
        return action.getName();
    }

    @Override
    public boolean hasController() {
        return action.hasController();
    }

    @Override
    public Object getController() {
        return action.getController();
    }

    @Override
    public Class<?> getControllerClass() {
        return action.getControllerClass();
    }

    @Override
    public ReflectMethod getMethod() {
        return action.getMethod();
    }

    @Override
    public boolean hasReturnValue() {
        return action.hasReturnValue();
    }

    @Override
    public boolean hasArguments() {
        return action.hasArguments();
    }

    @Override
    public Class<?> getReturnType() {
        return action.getReturnType();
    }

    @Override
    public Type getGenericReturnType() {
        return action.getGenericReturnType();
    }

    @Override
    public Argument[] getArguments() {
        return action.getArguments();
    }

    @Override
    public Annotation[] getAnnotations() {
        return action.getAnnotations();
    }

    @Override
    public <T extends Annotation> T[] getAnnotationsByType(Class<T> annotationClass) {
        return action.getAnnotationsByType(annotationClass);
    }

    @Override
    public Annotation[] getControllerAnnotations() {
        return action.getControllerAnnotations();
    }

    @Override
    public Annotation[] getMergedAnnotations() {
        return action.getMergedAnnotations();
    }

    @Override
    public ActionInterceptor[] getInterceptors() {
        return action.getInterceptors();
    }

    @Override
    public <T extends Annotation> T getControllerAnnotation(Class<T> annotationType) {
        return action.getControllerAnnotation(annotationType);
    }

    @Override
    public <T extends Annotation> T searchAnnotation(Class<T> annotationType) {
        return action.searchAnnotation(annotationType);
    }

    @Override
    public <T> void setExtension(T extension) {
        action.setExtension(extension);
    }

    @Override
    public <T extends Annotation> T getAnnotation(Class<T> annotationType) {
        return action.getAnnotation(annotationType);
    }

    @Override
    public boolean isAnnotationPresent(Class<? extends Annotation> annotationType) {
        return action.isAnnotationPresent(annotationType);
    }

    @Override
    public String toString() {
        return action.toString();
    }
}


text:
namespace UdpKit {
    struct UdpHeader {
        public const int SEQ_BITS = 15;
        public const int SEQ_PADD = 16 - SEQ_BITS;
        public const int SEQ_MASK = (1 << SEQ_BITS) - 1;
        public const int NETPING_BITS = 16;

        public ushort ObjSequence;
        public ushort AckSequence;
        public ulong AckHistory;
        public ushort AckTime;
        public bool IsObject;
        public uint Now;

        public void Pack (UdpStream buffer, UdpSocket socket) {
            int pos = buffer.Position;

            buffer.Position = 0;
            buffer.WriteUShort(PadSequence(ObjSequence), SEQ_BITS + SEQ_PADD);
            buffer.WriteUShort(PadSequence(AckSequence), SEQ_BITS + SEQ_PADD);
            buffer.WriteULong(AckHistory, UdpSocket.AckRedundancy);

            if (UdpSocket.CalculateNetworkPing) {
                buffer.WriteUShort(AckTime, NETPING_BITS);
            }

            buffer.Position = pos;
        }

        public void Unpack (UdpStream buffer, UdpSocket socket) {
            buffer.Position = 0;

            ObjSequence = TrimSequence(buffer.ReadUShort(SEQ_BITS + SEQ_PADD));
            AckSequence = TrimSequence(buffer.ReadUShort(SEQ_BITS + SEQ_PADD));
            AckHistory = buffer.ReadULong(UdpSocket.AckRedundancy);

            if (UdpSocket.CalculateNetworkPing) {
                AckTime = buffer.ReadUShort(NETPING_BITS);
            }
        }

        ushort PadSequence (ushort sequence) {
            sequence <<= SEQ_PADD;

            if (IsObject)
                sequence |= ((1 << SEQ_PADD) - 1);

            return sequence;
        }

        ushort TrimSequence (ushort sequence) {
            sequence >>= SEQ_PADD;
            return sequence;
        }
    }
}


text:
package com.mymobileapi.api5;

import javax.xml.bind.annotation.XmlAccessType;
import javax.xml.bind.annotation.XmlAccessorType;
import javax.xml.bind.annotation.XmlElement;
import javax.xml.bind.annotation.XmlRootElement;
import javax.xml.bind.annotation.XmlType;

/**
 *
Java class for anonymous complex type.
 *
 *
The following schema fragment specifies the expected content contained within this class.
 *
 *

 * <complexType>
 *   <complexContent>
 *     <restriction base="{http://www.w3.org/2001/XMLSchema}anyType">
 *       <sequence>
 *         <element name="Doc_DS_ZIPResult" type="{http://www.w3.org/2001/XMLSchema}base64Binary" minOccurs="0"/>
 *       </sequence>
 *     </restriction>
 *   </complexContent>
 * </complexType>
 *

 *
 *
 */
@XmlAccessorType(XmlAccessType.FIELD)
@XmlType(name = "", propOrder = {
    "docDSZIPResult"
})
@XmlRootElement(name = "Doc_DS_ZIPResponse")
public class DocDSZIPResponse {

    @XmlElement(name = "Doc_DS_ZIPResult")
    protected byte[] docDSZIPResult;

    /**
     * Gets the value of the docDSZIPResult property.
     *
     * @return
     *     possible object is
     *     byte[]
     */
    public byte[] getDocDSZIPResult() {
        return docDSZIPResult;
    }

    /**
     * Sets the value of the docDSZIPResult property.
     *
     * @param value
     *     allowed object is
     *     byte[]
     */
    public void setDocDSZIPResult(byte[] value) {
        this.docDSZIPResult = value;
    }

}


text:
package org.test.news;

import java.util.Arrays;
import java.util.HashMap;
import java.util.Map;
import java.util.Objects;

/**
 * Scores a series of space-separated words, based on the number of positive and negative words.
 *
 * @author Horia Chiorean (horia.chiorean@gmail.com)
 * @see Dictionary
 */
public class PhraseScorer {

    protected static final PhraseScorer INSTANCE = new PhraseScorer();

    private static final Map<String, Integer> SCORE_MAP;

    static {
        SCORE_MAP = new HashMap<>(Dictionary.INSTANCE.allWords().size());
        Dictionary.INSTANCE.goodWords().forEach(word -> SCORE_MAP.put(word, 1));
        Dictionary.INSTANCE.badWords().forEach(word -> SCORE_MAP.put(word, -1));
    }

    private PhraseScorer() {
    }

    protected int scoreMessage(String message) {
        Objects.requireNonNull(message);
        return Arrays.stream(message.split(" "))
                     .mapToInt(word -> SCORE_MAP.getOrDefault(word, 0))
                     .sum();
    }

    public static void main(String[] args) {
        //junit...sigh
        if (1 != INSTANCE.scoreMessage("up rise down")) {
            throw new IllegalStateException("incorrect");
        } else if (3 != INSTANCE.scoreMessage("up up up")) {
            throw new IllegalStateException("incorrect");
        } else if (0 != INSTANCE.scoreMessage("???")) {
            throw new IllegalStateException("incorrect");
        } else if (-1 != INSTANCE.scoreMessage("über fall bad")) {
            throw new IllegalStateException("incorrect");
        }
    }
}


text:
package org.apache.commons.functor.core.composite;

import org.apache.commons.functor.BinaryPredicate;

/**
 * {@link #test Tests} true iff
 * none of its children test false.
 * Note that by this definition, the "and" of
 * an empty collection of predicates tests true.
 *

 * Note that although this class implements
 * {@link java.io.Serializable Serializable}, a given instance will
 * only be truly Serializable if all the
 * underlying functors are.  Attempts to serialize
 * an instance whose delegates are not all
 * Serializable will result in an exception.
 *

 * @param <L> the left argument type.
 * @param <R> the right argument type.
 * @version $Revision: 1345136 $ $Date: 2012-06-01 09:47:06 -0300 (Fri, 01 Jun 2012) $
 */
public final class BinaryAnd<L, R> extends BaseBinaryPredicateList<L, R> {

    /**
     * serialVersionUID declaration.
     */
    private static final long serialVersionUID = -6741089498337923680L;

    // constructor
    // ------------------------------------------------------------------------
    /**
     * Create a new BinaryAnd.
     */
    public BinaryAnd() {
        super();
    }

    /**
     * Create a new BinaryAnd instance.
     *
     * @param predicates the predicates to add
     */
    public BinaryAnd(BinaryPredicate<? super L, ? super R>... predicates) {
        super(predicates);
    }

    /**
     * Create a new BinaryAnd instance.
     *
     * @param predicates the predicates to add
     */
    public BinaryAnd(Iterable<BinaryPredicate<? super L, ? super R>> predicates) {
        super(predicates);
    }

    // modifiers
    // ------------------------------------------------------------------------
    /**
     * And in a BinaryPredicate.
     * @param p BinaryPredicate to add
     * @return this
     */
    public BinaryAnd<L, R> and(BinaryPredicate<? super L, ? super R> p) {
        super.addBinaryPredicate(p);
        return this;
    }

    // predicate interface
    // ------------------------------------------------------------------------
    /**
     * {@inheritDoc}
     */
    public boolean test(L a, R b) {
        for (BinaryPredicate<? super L, ? super R> p : getBinaryPredicateList()) {
            if (!p.test(a, b)) {
                return false;
            }
        }
        return true;
    }

    /**
     * {@inheritDoc}
     */
    @Override
    public boolean equals(Object that) {
        return that == this || (that instanceof BinaryAnd<?, ?> && equals((BinaryAnd<?, ?>) that));
    }

    /**
     * Learn whether another BinaryAnd is equal to this.
     * @param that the BinaryAnd to test
     * @return boolean
     */
    public boolean equals(BinaryAnd<?, ?> that) {
        return getBinaryPredicateListEquals(that);
    }

    /**
     * {@inheritDoc}
     */
    @Override
    public int hashCode() {
        return "BinaryAnd".hashCode() ^ getBinaryPredicateListHashCode();
    }

    /**
     * {@inheritDoc}
     */
    @Override
    public String toString() {
        return "BinaryAnd<" + getBinaryPredicateListToString() + ">";
    }

}


text:
#include <vxGL/ProgramPipeline.h>
#include <vxGL/ShaderProgram.h>
#include <memory>
#include <fstream>
#include <vxGL/gl.h>

namespace vx
{
	namespace gl
	{
		namespace detail
		{
			inline u32 getStageBit(ShaderProgramType type)
			{
				switch (type)
				{
				case gl::ShaderProgramType::VERTEX:
					return GL_VERTEX_SHADER_BIT;
				case gl::ShaderProgramType::GEOMETRY:
					return GL_GEOMETRY_SHADER_BIT;
				case gl::ShaderProgramType::COMPUTE:
					return GL_COMPUTE_SHADER_BIT;
				case gl::ShaderProgramType::FRAGMENT:
					return GL_FRAGMENT_SHADER_BIT;
				case gl::ShaderProgramType::TESS_CONTROL:
					return GL_TESS_CONTROL_SHADER_BIT;
				case gl::ShaderProgramType::TESS_EVAL:
					return GL_TESS_EVALUATION_SHADER_BIT;
				default:
					return 0;
				}
			}
		}

		std::unique_ptr<char[]> loadProgramData(const char *filename)
		{
			std::ifstream inFile(filename);
			if (!inFile.is_open())
				return nullptr;

			inFile.seekg(0, inFile.end);
			u32 length = inFile.tellg();
			inFile.seekg(0, inFile.beg);

			auto ptr = std::unique_ptr<char[]>(new char[length + 1]);

			inFile.read(ptr.get(), length);

			ptr[length] = '\0';

			//if (pCount)
			//	*pCount = length + 1;

			return std::move(ptr);
		}

		ProgramPipeline::ProgramPipeline() noexcept
			:Base(),
			m_programIds()
		{
		}

		ProgramPipeline::ProgramPipeline(ProgramPipeline &&rhs) noexcept
			: Base(std::move(rhs))
		{
			memcpy(m_programIds, rhs.m_programIds, sizeof(m_programIds));
		}

		ProgramPipeline::~ProgramPipeline()
		{
			destroy();
		}

		ProgramPipeline& ProgramPipeline::operator=(ProgramPipeline &&rhs) noexcept
		{
			if (this != &rhs)
			{
				Base::operator=(std::move(rhs));
				std::swap(m_programIds, rhs.m_programIds);
			}
			return *this;
		}

		void ProgramPipeline::create()
		{
			if (m_id == 0)
			{
				glCreateProgramPipelines(1, &m_id);
			}
		}

		void ProgramPipeline::destroy()
		{
			if (m_id != 0)
			{
				glDeleteProgramPipelines(1, &m_id);
				m_id = 0;
			}
		}

		void ProgramPipeline::useProgram(const ShaderProgram &program)
		{
			auto programId = program.getId();
			auto type = program.getType();
			glUseProgramStages(m_id, ::vx::gl::detail::getStageBit(type), programId);

			m_programIds[static_cast<u32>(type)] = programId;
		}

		void ProgramPipeline::validatePipeline() const
		{
			glValidateProgramPipeline(m_id);
		}

		void ProgramPipeline::bind() const
		{
			glBindProgramPipeline(m_id);
		}

		void ProgramPipeline::bindZero() const
		{
			glBindProgramPipeline(0);
		}

		u32 ProgramPipeline::operator[](ShaderProgramType type) const
		{
			return m_programIds[static_cast<u32>(type)];
		}

		u32 ProgramPipeline::getVertexShader() const noexcept
		{
			return m_programIds[static_cast<u32>(ShaderProgramType::VERTEX)];
		}

		u32 ProgramPipeline::getFragmentShader() const noexcept
		{
			return m_programIds[static_cast<u32>(ShaderProgramType::FRAGMENT)];
		}

		u32 ProgramPipeline::getGeometryShader() const noexcept
		{
			return m_programIds[static_cast<u32>(ShaderProgramType::GEOMETRY)];
		}

			u32 ProgramPipeline::getComputeShader() const noexcept
		{
			return m_programIds[static_cast<u32>(ShaderProgramType::COMPUTE)];
		}

		u32 ProgramPipeline::getTessControlShader() const noexcept
		{
			return m_programIds[static_cast<u32>(ShaderProgramType::TESS_CONTROL)];
		}

		u32 ProgramPipeline::getTessEvalShader() const noexcept
		{
			return m_programIds[static_cast<u32>(ShaderProgramType::TESS_EVAL)];
		}
	}
}