TestPager.kt

/*
 * Copyright 2022 The Android Open Source Project
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package androidx.paging.testing

import androidx.paging.LoadType
import androidx.paging.LoadType.APPEND
import androidx.paging.LoadType.PREPEND
import androidx.paging.LoadType.REFRESH
import androidx.paging.PagingConfig
import androidx.paging.PagingSource
import androidx.paging.PagingSource.LoadParams
import androidx.paging.PagingSource.LoadResult
import androidx.paging.Pager
import androidx.paging.PagingSource.LoadResult.Page.Companion.COUNT_UNDEFINED
import androidx.paging.PagingState
import java.util.concurrent.atomic.AtomicBoolean
import kotlinx.coroutines.sync.Mutex
import kotlinx.coroutines.sync.withLock

/**
 * A fake [Pager] class to simulate how a real Pager and UI would load data from a PagingSource.
 *
 * As Paging's first load is always of type [LoadType.REFRESH], the first load operation of
 * the [TestPager] must be a call to [refresh].
 *
 * This class only supports loads from a single instance of PagingSource. To simulate
 * multi-generational Paging behavior, you must create a new [TestPager] by supplying a
 * new instance of [PagingSource].
 *
 * @param pagingSource the [PagingSource] to load data from.
 * @param config the [PagingConfig] to configure this TestPager's loading behavior.
 */
public class TestPager<Key : Any, Value : Any>(
    private val pagingSource: PagingSource<Key, Value>,
    private val config: PagingConfig,
) {
    private val hasRefreshed = AtomicBoolean(false)

    private val lock = Mutex()

    private val pages = ArrayDeque<LoadResult.Page<Key, Value>>()

    /**
     * Performs a load of [LoadType.REFRESH] on the PagingSource.
     *
     * If initialKey != null, refresh will start loading from the supplied key.
     *
     * Since Paging's first load is always of [LoadType.REFRESH], this method must be the very
     * first load operation to be called on the TestPager before either [append] or [prepend]
     * can be called. However, other non-loading operations can still be invoked. For example,
     * you can call [getLastLoadedPage] before any load operations.
     *
     * Returns the LoadResult upon refresh on the [PagingSource].
     *
     * @param initialKey the [Key] to start loading data from on initial refresh.
     *
     * @throws IllegalStateException TestPager does not support multi-generational paging behavior.
     * As such, multiple calls to refresh() on this TestPager is illegal. The [PagingSource] passed
     * in to this [TestPager] will also be invalidated to prevent reuse of this pager for loads.
     * However, other [TestPager] methods that does not invoke loads can still be called,
     * such as [getLastLoadedPage].
     */
    public suspend fun refresh(
        initialKey: Key? = null
    ): @JvmSuppressWildcards LoadResult<Key, Value> {
        if (!hasRefreshed.compareAndSet(false, true)) {
            pagingSource.invalidate()
            throw IllegalStateException("TestPager does not support multi-generational access " +
                "and refresh() can only be called once per TestPager. To start a new generation," +
                "create a new TestPager with a new PagingSource.")
        }
        return doInitialLoad(initialKey)
    }

    /**
     * Performs a load of [LoadType.APPEND] on the PagingSource.
     *
     * Since Paging's first load is always of [LoadType.REFRESH], [refresh] must always be called
     * first before this [append] is called.
     *
     * If [PagingConfig.maxSize] is implemented, [append] loads that exceed [PagingConfig.maxSize]
     * will cause pages to be dropped from the front of loaded pages.
     *
     * Returns the [LoadResult] from calling [PagingSource.load]. If the [LoadParams.key] is null,
     * such as when there is no more data to append, this append will be no-op by returning null.
     */
    public suspend fun append(): @JvmSuppressWildcards LoadResult<Key, Value>? {
        return doLoad(APPEND)
    }

    /**
     * Performs a load of [LoadType.PREPEND] on the PagingSource.
     *
     * Since Paging's first load is always of [LoadType.REFRESH], [refresh] must always be called
     * first before this [prepend] is called.
     *
     * If [PagingConfig.maxSize] is implemented, [prepend] loads that exceed [PagingConfig.maxSize]
     * will cause pages to be dropped from the end of loaded pages.
     *
     * Returns the [LoadResult] from calling [PagingSource.load]. If the [LoadParams.key] is null,
     * such as when there is no more data to prepend, this prepend will be no-op by returning null.
     */
    public suspend fun prepend(): @JvmSuppressWildcards LoadResult<Key, Value>? {
        return doLoad(PREPEND)
    }

    /**
     * Helper to perform REFRESH loads.
     */
    private suspend fun doInitialLoad(
        initialKey: Key?
    ): @JvmSuppressWildcards LoadResult<Key, Value> {
        return lock.withLock {
            pagingSource.load(
                LoadParams.Refresh(initialKey, config.initialLoadSize, config.enablePlaceholders)
            ).also { result ->
                if (result is LoadResult.Page) {
                    pages.addLast(result)
                }
            }
        }
    }

    /**
     * Helper to perform APPEND or PREPEND loads.
     */
    private suspend fun doLoad(loadType: LoadType): LoadResult<Key, Value>? {
        return lock.withLock {
            if (!hasRefreshed.get()) {
                throw IllegalStateException("TestPager's first load operation must be a refresh. " +
                    "Please call refresh() once before calling ${loadType.name.lowercase()}().")
            }
            when (loadType) {
                REFRESH -> throw IllegalArgumentException(
                    "For LoadType.REFRESH use doInitialLoad()"
                )
                APPEND -> {
                    val key = pages.lastOrNull()?.nextKey ?: return null
                    pagingSource.load(
                        LoadParams.Append(key, config.pageSize, config.enablePlaceholders)
                    ).also { result ->
                        if (result is LoadResult.Page) {
                            pages.addLast(result)
                        }
                        dropPagesOrNoOp(PREPEND)
                    }
                } PREPEND -> {
                    val key = pages.firstOrNull()?.prevKey ?: return null
                    pagingSource.load(
                        LoadParams.Prepend(key, config.pageSize, config.enablePlaceholders)
                    ).also { result ->
                        if (result is LoadResult.Page) {
                            pages.addFirst(result)
                        }
                        dropPagesOrNoOp(APPEND)
                    }
                }
            }
        }
    }

    /**
     * Returns the most recent [LoadResult.Page] loaded from the [PagingSource]. Null if
     * no pages have been returned from [PagingSource]. For example, if PagingSource has
     * only returned [LoadResult.Error] or [LoadResult.Invalid].
     */
    public suspend fun getLastLoadedPage(): @JvmSuppressWildcards LoadResult.Page<Key, Value>? {
        return lock.withLock {
            pages.lastOrNull()
        }
    }

    /**
     * Returns the current list of [LoadResult.Page] loaded so far from the [PagingSource].
     */
    public suspend fun getPages(): @JvmSuppressWildcards List<LoadResult.Page<Key, Value>> {
        return lock.withLock {
            pages.toList()
        }
    }

    /**
     * Returns a [PagingState] to generate a [LoadParams.key] by supplying it to
     * [PagingSource.getRefreshKey]. The key returned from [PagingSource.getRefreshKey]
     * should be used as the [LoadParams.Refresh.key] when calling [refresh] on a new generation of
     * TestPager.
     *
     * The anchorPosition must be within index of loaded items, which can include
     * placeholders if [PagingConfig.enablePlaceholders] is true. For example:
     * - No placeholders: If 40 items have been loaded so far , anchorPosition must be
     * in [0 .. 39].
     * - With placeholders: If there are a total of 100 loadable items, the anchorPosition
     * must be in [0..99].
     *
     * The [anchorPosition] should be the index that the user has hypothetically
     * scrolled to on the UI. Since the [PagingState.anchorPosition] in Paging can be based
     * on any item or placeholder currently visible on the screen, the actual
     * value of [PagingState.anchorPosition] may not exactly match the [anchorPosition] passed
     * to this function even if viewing the same page of data.
     *
     * Note that when `[PagingConfig.enablePlaceholders] = false`, the
     * [PagingState.anchorPosition] returned from this function references the absolute index
     * within all loadable data. For example, with items[0 - 99]:
     * If items[20 - 30] were loaded without placeholders, anchorPosition 0 references item[20].
     * But once translated into [PagingState.anchorPosition], anchorPosition 0 references item[0].
     * The [PagingSource] is expected to handle this correctly within [PagingSource.getRefreshKey]
     * when [PagingConfig.enablePlaceholders] = false.
     *
     * @param anchorPosition the index representing the last accessed item within the
     * items presented on the UI, which may be a placeholder if
     * [PagingConfig.enablePlaceholders] is true.
     *
     * @throws IllegalStateException if anchorPosition is out of bounds.
     */
    public suspend fun getPagingState(
        anchorPosition: Int
    ): @JvmSuppressWildcards PagingState<Key, Value> {
        lock.withLock {
            checkWithinBoundary(anchorPosition)
            return PagingState(
                pages = pages.toList(),
                anchorPosition = anchorPosition,
                config = config,
                leadingPlaceholderCount = getLeadingPlaceholderCount()
            )
        }
    }

    /**
     * Returns a [PagingState] to generate a [LoadParams.key] by supplying it to
     * [PagingSource.getRefreshKey]. The key returned from [PagingSource.getRefreshKey]
     * should be used as the [LoadParams.Refresh.key] when calling [refresh] on a new generation of
     * TestPager.
     *
     * The [anchorPositionLookup] lambda should return an item that the user has hypothetically
     * scrolled to on the UI. The item must have already been loaded prior to using this helper.
     * To generate a PagingState anchored to a placeholder, use the overloaded [getPagingState]
     * function instead. Since the [PagingState.anchorPosition] in Paging can be based
     * on any item or placeholder currently visible on the screen, the actual
     * value of [PagingState.anchorPosition] may not exactly match the anchorPosition returned
     * from this function even if viewing the same page of data.
     *
     * Note that when `[PagingConfig.enablePlaceholders] = false`, the
     * [PagingState.anchorPosition] returned from this function references the absolute index
     * within all loadable data. For example, with items[0 - 99]:
     * If items[20 - 30] were loaded without placeholders, anchorPosition 0 references item[20].
     * But once translated into [PagingState.anchorPosition], anchorPosition 0 references item[0].
     * The [PagingSource] is expected to handle this correctly within [PagingSource.getRefreshKey]
     * when [PagingConfig.enablePlaceholders] = false.
     *
     * @param anchorPositionLookup the predicate to match with an item which will serve as the basis
     * for generating the [PagingState].
     *
     * @throws IllegalArgumentException if the given predicate fails to match with an item.
     */
    public suspend fun getPagingState(
        anchorPositionLookup: (item: @JvmSuppressWildcards Value) -> Boolean
    ): @JvmSuppressWildcards PagingState<Key, Value> {
        lock.withLock {
            val indexInPages = pages.flatten().indexOfFirst {
                anchorPositionLookup(it)
            }
            return when {
                indexInPages < 0 -> throw IllegalArgumentException(
                    "The given predicate has returned false for every loaded item. To generate a" +
                        "PagingState anchored to an item, the expected item must have already " +
                        "been loaded."
                )
                else -> {
                    val finalIndex = if (config.enablePlaceholders) {
                        indexInPages + (pages.firstOrNull()?.itemsBefore ?: 0)
                    } else {
                        indexInPages
                    }
                    PagingState(
                        pages = pages.toList(),
                        anchorPosition = finalIndex,
                        config = config,
                        leadingPlaceholderCount = getLeadingPlaceholderCount()
                    )
                }
            }
        }
    }

    /**
     * Ensures the anchorPosition is within boundary of loaded data.
     *
     * If placeholders are enabled, the provided anchorPosition must be within boundaries of
     * [0 .. itemCount - 1], which includes placeholders before and after loaded data.
     *
     * If placeholders are disabled, the provided anchorPosition must be within boundaries of
     * [0 .. loaded data size - 1].
     *
     * @throws IllegalStateException if anchorPosition is out of bounds
     */
    private fun checkWithinBoundary(anchorPosition: Int) {
        val loadedSize = pages.flatten().size
        val maxBoundary = if (config.enablePlaceholders) {
            (pages.firstOrNull()?.itemsBefore ?: 0) + loadedSize +
                (pages.lastOrNull()?.itemsAfter ?: 0) - 1
        } else {
            loadedSize - 1
        }
        check(anchorPosition in 0..maxBoundary) {
            "anchorPosition $anchorPosition is out of bounds between [0..$maxBoundary]. Please " +
                "provide a valid anchorPosition."
        }
    }

    // Number of placeholders before the first loaded item if placeholders are enabled, otherwise 0.
    private fun getLeadingPlaceholderCount(): Int {
        return if (config.enablePlaceholders) {
            // itemsBefore represents placeholders before first loaded item, and can be
            // one of three.
            // 1. valid int if implemented
            // 2. null if pages empty
            // 3. COUNT_UNDEFINED if not implemented
            val itemsBefore: Int? = pages.firstOrNull()?.itemsBefore
            // finalItemsBefore is `null` if it is either case 2. or 3.
            val finalItemsBefore = if (itemsBefore == null || itemsBefore == COUNT_UNDEFINED) {
                null
            } else {
                itemsBefore
            }
            // This will ultimately return 0 if user didn't implement itemsBefore or if pages
            // are empty, i.e. user called getPagingState before any loads.
            finalItemsBefore ?: 0
        } else {
            0
        }
    }

    private fun dropPagesOrNoOp(dropType: LoadType) {
        require(dropType != REFRESH) {
            "Drop loadType must be APPEND or PREPEND but got $dropType"
        }

        // check if maxSize has been set
        if (config.maxSize == PagingConfig.MAX_SIZE_UNBOUNDED) return

        var itemCount = pages.flatten().size
        if (itemCount < config.maxSize) return

        // represents the max droppable amount of items
        val presentedItemsBeforeOrAfter = when (dropType) {
            PREPEND -> pages.take(pages.lastIndex)
            else -> pages.takeLast(pages.lastIndex)
        }.fold(0) { acc, page ->
            acc + page.data.size
        }

        var itemsDropped = 0

        // mirror Paging requirement to never drop below 2 pages
        while (pages.size > 2 && itemCount - itemsDropped > config.maxSize) {
            val pageSize = when (dropType) {
                PREPEND -> pages.first().data.size
                else -> pages.last().data.size
            }

            val itemsAfterDrop = presentedItemsBeforeOrAfter - itemsDropped - pageSize

            // mirror Paging behavior of ensuring prefetchDistance is fulfilled in dropped
            // direction
            if (itemsAfterDrop < config.prefetchDistance) break

            when (dropType) {
                PREPEND -> pages.removeFirst()
                else -> pages.removeLast()
            }

            itemsDropped += pageSize
        }
    }
}