--[[----------------------------------------------------------------------------

WFRxRateLimiterTest.lua

--------------------------------------------------------------------------------

ADOBE CONFIDENTIAL
------------------
Copyright 2016 Adobe Systems Incorporated
All Rights Reserved.

NOTICE: All information contained herein is, and remain the property of
Adobe Systems Incorporated and its suppliers, if any. The intellectual and
technical concepts contained herein are proprietary to Adobe Systems
Incorporated and its suppliers and may be covered by U.S. and Foreign
Patents, patents in process, and are protected by trade secret or copyright
law.  Dissemination of this information or reproduction of this material is
strictly forbidden unless prior written permission is obtained from
Adobe Systems Incorporated.

------------------------------------------------------------------------------]]

	-- RxLua
local Rx = require 'Rx'
local Observable = Rx.Observable
local WFRxScheduler = require 'WFRxScheduler'

	-- wichitafoundation.core.self_test
local WFRxTestCase = require 'WFRxTestCase'

	-- wichitafoundation.core.utils
local WFRxRateLimiter = require 'WFRxRateLimiter'
local WFQueueing = require 'WFQueueing'
local WFRxRunUntilDone = require 'WFRxRunUntilDone'

--------------------------------------------------------------------------------

local WFRxRateLimiterTest = {}

local runUntilDone = WFRxRunUntilDone

--------------------------------------------------------------------------------

function WFRxRateLimiterTest.rxTestNamespace( rxTest )
	assert( WFRxRateLimiter.makeRateLimiter, "No makeRateLimiter routine" )
end

--------------------------------------------------------------------------------

function WFRxRateLimiterTest.rxTestMakeRateLimiter( rxTest )

    local kSimultaneousTasksLimit = 5

    local limiter = WFRxRateLimiter.makeRateLimiter(
            WFQueueing.makePriorityQueue,
            kSimultaneousTasksLimit )

    local limitedObservables = {}

    local numberOfTasks = 20

    -- Create low priority tasks

    local maxRunning = 0
	local curRunning = 0
	 
    for x = 1, numberOfTasks do

        local simpleObservable = Observable.just( "LOW" )
            :tap( function()
                curRunning = curRunning + 1
                maxRunning = math.max( curRunning, maxRunning )
            end )
            :delay( 0.2, WFRxScheduler.sharedScheduler() )
            :tap( function()
                curRunning = curRunning - 1
            end )

        limitedObservables[ x ] = limiter:enqueueObservable( simpleObservable, 0 ):call( runUntilDone )

    end

    -- Let the low priority tasks start.

    rxTest:sleep( 0.05 )

	rxTest:expectEqual( curRunning, maxRunning )

    local merged = limitedObservables[ 1 ] 

    for x = 2, numberOfTasks do
        merged = merged:merge( limitedObservables[ x ] )
    end

    local mergedObserver = rxTest:observe( merged )

    for x = 1, numberOfTasks do
        mergedObserver:expectData( "LOW" )
    end

    rxTest:expectEqual( kSimultaneousTasksLimit, maxRunning )

end

--------------------------------------------------------------------------------

function WFRxRateLimiterTest.rxTestSimpleEnqueueObservableAndPriority( rxTest )

	rxTest:setTimeLimit( 10 )

	local kNumTasksToRun = 20
	local kSimultaneousTasksLimit = 1

	local numStarted = 0
	local lastPriority = math.huge
	
	local limiter = WFRxRateLimiter.makeRateLimiter(
			WFQueueing.makePriorityQueue,
			kSimultaneousTasksLimit )

	local limitedObservables = {}

	local allTasksCompleted = Rx.Subject()

	for x = 1, kNumTasksToRun do
	
		local priority = math.random( 10 )
	
		local delay = 0
		
		if x == 1 then
			delay = 1
			-- Because the rate limiter will allow a higher priority
			-- task to bypass the targetSimultaneousCost we need to
			-- default this first task to be the highest priority, the
			-- enqueued results will then be issued in priority order
			-- so it won't issue more than one task at a time just as before.
			priority = math.huge
		end
	
		local simpleObservable = Observable.timer( delay, WFRxScheduler.sharedScheduler() )
			:map( function()
			
				if numStarted > 1 then
					assert( priority <= lastPriority )
				end
				
				lastPriority = priority
				numStarted = numStarted + 1
				
				return numStarted
				
			end )
			:tap( function( numStarted )
			
				if numStarted == kNumTasksToRun then
					allTasksCompleted:onCompleted()
				end
			
			end )
			 
		limitedObservables[ x ] = limiter:enqueueObservable( simpleObservable, priority )
		
	end
	
	rxTest:sleep( 0.5 )
	
	rxTest:expectEqual( 0, numStarted )
	
	-- The requests aren't actually queued up until the underlying observable is subscribed to.
	-- Once a request is queued, it will be scheduled to execute once their is a slot
	-- available. This means the very first request we schedule here will execute
	-- immediately since the queue is otherwise empty. Even if the priority of that first
	-- request is low, it will run right away since the queue is empty. However, we can
	-- safely test the priorities of all but the first request since our
	-- kSimultaneousTasksLimit is 1, and all the requests should be subscribed to before
	-- the first request finishes.
	
	for x = 1, kNumTasksToRun do
		limitedObservables[ x ]:call( runUntilDone )
	end

	rxTest:expectOnlyCompleted( allTasksCompleted )

end

--------------------------------------------------------------------------------

function WFRxRateLimiterTest.rxTestHighPrioritySimple( rxTest )

	rxTest:setTimeLimit( 10 )

	local kSimultaneousTasksLimit = 1

	local limiter = WFRxRateLimiter.makeRateLimiter(
			WFQueueing.makePriorityQueue,
			kSimultaneousTasksLimit )

	local highPriorityRequestsCompleted = 0
	local lowPriorityRequestsCompleted = 0
	
	-- Create low priority task

	local simpleObservable = Observable.just( "LOW" )
	   :delay( 0.5, WFRxScheduler.sharedScheduler() )
	   :tap( function() lowPriorityRequestsCompleted = lowPriorityRequestsCompleted + 1 end )

	local lowPriorityTask = limiter:enqueueObservable( simpleObservable, 0 ):call( runUntilDone )

	-- Let the low priority tasks start.
	
	rxTest:sleep( 0.05 )

	-- Now start high priority tasks. The low priority tasks should
	-- be cancelled and rescheduled after the high priority tasks complete.

	local simpleObservable2 = Observable.just( "HIGH" )
	   :delay( 0.2, WFRxScheduler.sharedScheduler() )
	   :tap( function() highPriorityRequestsCompleted = highPriorityRequestsCompleted + 1 end )
		
	local highPriorityTask = limiter:enqueueObservable( simpleObservable2, 10 ):call( runUntilDone )

	
	-- Verify the high priority observables complete before the low prioirty observables.

	local merged = lowPriorityTask:merge( highPriorityTask )
	
	local mergedObserver = rxTest:observe( merged )
	
	mergedObserver:expectData( "HIGH" )
	rxTest:expectEqual( highPriorityRequestsCompleted, 1 )
	rxTest:expectEqual( lowPriorityRequestsCompleted, 0 )
	
	mergedObserver:expectData( "LOW" )
	rxTest:expectEqual( highPriorityRequestsCompleted, 1 )
	rxTest:expectEqual( lowPriorityRequestsCompleted, 1 )

	mergedObserver:expectCompleted()
	
end

--------------------------------------------------------------------------------

function WFRxRateLimiterTest.rxTestHighPriorityObservableMultipleTimes( rxTest )

	rxTest:setTimeLimit( 10 )

	local kSimultaneousTasksLimit = 1

	local limiter = WFRxRateLimiter.makeRateLimiter(
			WFQueueing.makePriorityQueue,
			kSimultaneousTasksLimit )

	local highPriorityRequestsCompleted = 0
	local lowPriorityRequestsCompleted = 0
	
	-- Create low priority task

	local simpleObservable = Observable.just( "LOW" )
	   :delay( 0.6, WFRxScheduler.sharedScheduler() )
	   :tap( function() lowPriorityRequestsCompleted = lowPriorityRequestsCompleted + 1 end )

	local lowPriorityTask = limiter:enqueueObservable( simpleObservable, 0 ):call( runUntilDone )

	-- Let the low priority tasks start.
	
	rxTest:sleep( 0.05 )

	-- Now start high priority tasks. The high priority tasks complete
	-- should be started even though it exceeds the simultaneous task limit,
	-- and it should finish before the low priority task as it's delay is shorter.

	local simpleObservable2 = Observable.just( "HIGH" )
	   :delay( 0.1, WFRxScheduler.sharedScheduler() )
	   :tap( function() highPriorityRequestsCompleted = highPriorityRequestsCompleted + 1 end )
		
	local highPriorityTask = limiter:enqueueObservable( simpleObservable2, 10 ):call( runUntilDone )
	
	-- Let the High priority task finish
	
	local merged = lowPriorityTask:merge( highPriorityTask )
	local mergedObserver = rxTest:observe( merged )
	
	mergedObserver:expectData( "HIGH" )
	rxTest:expectEqual( highPriorityRequestsCompleted, 1 )
	rxTest:expectEqual( lowPriorityRequestsCompleted, 0 )
	
	-- Let the low priority task have more time.
	
	rxTest:sleep( 0.05 )
	
	-- Start another high priority task.
	
	highPriorityTask = limiter:enqueueObservable( simpleObservable2, 10 ):call( runUntilDone )

	merged = merged:merge( highPriorityTask )
	mergedObserver = rxTest:observe( merged )
	
	mergedObserver:expectData( "HIGH" )
	rxTest:expectEqual( highPriorityRequestsCompleted, 1 )
	rxTest:expectEqual( lowPriorityRequestsCompleted, 0 )
	
	mergedObserver:expectData( "HIGH" )
	rxTest:expectEqual( highPriorityRequestsCompleted, 2 )
	rxTest:expectEqual( lowPriorityRequestsCompleted, 0 )
	
	mergedObserver:expectData( "LOW" )
	rxTest:expectEqual( highPriorityRequestsCompleted, 2 )
	rxTest:expectEqual( lowPriorityRequestsCompleted, 1 )

	mergedObserver:expectCompleted()
	
end

--------------------------------------------------------------------------------

function WFRxRateLimiterTest.rxTestHighPriorityMultiple( rxTest )

	rxTest:setTimeLimit( 50 )

	local kSimultaneousTasksLimit = 5

	local limiter = WFRxRateLimiter.makeRateLimiter(
			WFQueueing.makePriorityQueue,
			kSimultaneousTasksLimit )

	local limitedObservables = {}

	local highPriorityRequestsCompleted = 0
	local lowPriorityRequestsCompleted = 0
	
	local numberOfTasks = 20
	
	-- Create low priority tasks
	
	for x = 1, numberOfTasks / 2 do
	
		local simpleObservable = Observable.just( "LOW" )
		   :delay( 1.5, WFRxScheduler.sharedScheduler() )
		   :tap( function() lowPriorityRequestsCompleted = lowPriorityRequestsCompleted + 1 end )

		limitedObservables[ x ] = limiter:enqueueObservable( simpleObservable, 0, 1 )
			:call( runUntilDone )
		
	end
	
	-- Let the low priority tasks start.
	
	rxTest:sleep( 0.05 )

	-- Now start high priority tasks. The low priority tasks should
	-- be cancelled and rescheduled after the high priority tasks complete.

	for x = numberOfTasks / 2 + 1, numberOfTasks do
	
		local simpleObservable = Observable.just( "HIGH" )
		   :delay( 0.001, WFRxScheduler.sharedScheduler() )
		   :tap( function() highPriorityRequestsCompleted = highPriorityRequestsCompleted + 1 end )
			
		limitedObservables[ x ] = limiter:enqueueObservable( simpleObservable, 10, 1 )
			:call( runUntilDone )
		
	end
	
	-- Verify the high priority observables complete before the low priority observables
	-- not because they were rescheduled but because they were high priority and could
	-- bypass the simultaneous connection limit and their duration was less than the
	-- low priority tasks.

	local merged = limitedObservables[ 1 ] 

	for x = 2, numberOfTasks do
		merged = merged:merge( limitedObservables[ x ] )
	end
	
	local mergedObserver = rxTest:observe( merged )
	
	for x = 1, numberOfTasks / 2 do
		mergedObserver:expectData( "HIGH" )
		rxTest:expectEqual( highPriorityRequestsCompleted, x )
		rxTest:expectEqual( lowPriorityRequestsCompleted, 0 )
	end

	for x = 1, numberOfTasks / 2 do
		mergedObserver:expectData( "LOW" )
		rxTest:expectEqual( lowPriorityRequestsCompleted, x )
	end
	
	mergedObserver:expectCompleted()

end

--------------------------------------------------------------------------------

function WFRxRateLimiterTest.rxTestDisposing( rxTest )

	rxTest:setTimeLimit( 5 )

	local kSimultaneousTasksLimit = 1

	local limiter = WFRxRateLimiter.makeRateLimiter(
			WFQueueing.makePriorityQueue,
			kSimultaneousTasksLimit )

	-- Create a task

	local gotHere = false

	local function startTask()
	
		local simpleObservable = Observable.timer( 0.5, WFRxScheduler.sharedScheduler() )
			:take( 1 )
			:tap( function()
				gotHere = true
			end )
	
		return limiter:enqueueObservable( simpleObservable, 5 ):call( runUntilDone )
		
	end
	
	local task = startTask()
	
	-- Let the task start.
	
	rxTest:sleep( 0.1 )

 	-- Now dispose of the task
	
	task:dispose()
	
	task = nil
	
	-- Sleep long enough for the observable to finish
	
	rxTest:sleep( 1 )

	rxTest:expectEqual( gotHere, false )

end

--------------------------------------------------------------------------------

function WFRxRateLimiterTest.rxTestDisposeOpensSlots( rxTest )

	rxTest:setTimeLimit( 5 )

	local kSimultaneousTasksLimit = 1

	local limiter = WFRxRateLimiter.makeRateLimiter(
			WFQueueing.makePriorityQueue,
			kSimultaneousTasksLimit )

	-- Create a task

	local gotHere1 = false
	local gotHere2 = false
	
	local function startTask( taskId )
	
		local simpleObservable = Observable.timer( 0.5, WFRxScheduler.sharedScheduler() )
			:take( 1 )
			:tap( function()
				if taskId == 1 then
					gotHere1 = true
				else
					gotHere2 = true
				end
			end )
	
		return limiter:enqueueObservable( simpleObservable, 5 ):call( runUntilDone )
		
	end
	
	local task1 = startTask( 1 )
	startTask( 2 )
	
	-- Let the task start.
	
	rxTest:sleep( 0.1 )

 	-- Now dispose of the task
	
	task1:dispose()
	task1 = nil
	
	-- Sleep long enough for the task 2 to finish
	
	rxTest:sleep( 2 )

	rxTest:expectEqual( gotHere1, false )
	rxTest:expectEqual( gotHere2, true )
	
end

--------------------------------------------------------------------------------

function WFRxRateLimiterTest.rxTestSimpleCost( rxTest )

	rxTest:setTimeLimit( 10 )

	local kNumTasksToRun = 20
	local kSimultaneousTasksLimit = 5
	
	local inFlightCost = 0
	local numStarted = 0
	
	local limiter = WFRxRateLimiter.makeRateLimiter(
			WFQueueing.makePriorityQueue,
			kSimultaneousTasksLimit )

	local limitedObservables = {}

	local allTasksCompleted = Rx.Subject()

	for x = 1, kNumTasksToRun do
	
		local priority = x % 3
		local cost = 1

		local delay = 0
		
		if x == 1 then
			delay = 1
		end
	
		local simpleObservable = Observable.timer( delay, WFRxScheduler.sharedScheduler() )
			:map( function()
			
				inFlightCost = inFlightCost + cost
				assert( inFlightCost <= kSimultaneousTasksLimit + 2 ) -- + 2 = 1 for priority 2 bypassing limit, and priority 3 bypassing limit again
				numStarted = numStarted + 1
				return numStarted
				
			end )
			:delay( 0.1, WFRxScheduler.sharedScheduler() )
			:map( function ( numStarted )
				inFlightCost = inFlightCost - cost
				return numStarted
			end )
			:tap( function( numStarted )
			
				if numStarted == kNumTasksToRun then
					allTasksCompleted:onCompleted()
				end
			
			end )
			 
		limitedObservables[ x ] = limiter:enqueueObservable( simpleObservable, priority, cost )
		
	end
	
	for x = 1, kNumTasksToRun do
		limitedObservables[ x ] = limitedObservables[ x ]:call( runUntilDone )
	end

	rxTest:expectOnlyCompleted( allTasksCompleted )

end

--------------------------------------------------------------------------------

function WFRxRateLimiterTest.rxTestSimpleTimestamp( rxTest )

	rxTest:setTimeLimit( 10 )

	local kNumTasksToRun = 6
	local kSimultaneousTasksLimit = 5
	
	local inFlightCost = 0
	local numStarted = 0
	
	local limiter = WFRxRateLimiter.makeRateLimiter(
			WFQueueing.makePriorityQueue,
			kSimultaneousTasksLimit )

	local limitedObservables = {}

	local allTasksCompleted = Rx.Subject()

	local function makeObservable( priority, cost ) 
	
		return Observable.timer( 0.1, WFRxScheduler.sharedScheduler() )
			:map( function()
			
				inFlightCost = inFlightCost + cost
				assert( inFlightCost <= kSimultaneousTasksLimit )
				numStarted = numStarted + 1
				return numStarted
				
			end )
			:delay( 0.1, WFRxScheduler.sharedScheduler() )
			:map( function ( numStarted )
				inFlightCost = inFlightCost - cost
				return numStarted
			end )
			:tap( function( numStarted )
			
				if numStarted == kNumTasksToRun then
					allTasksCompleted:onCompleted()
				end
			
			end )
		
	end

	-- Create enough observables to fill all slots in the rate limiter.

	for x = 1, kSimultaneousTasksLimit do

		local priority = 5
		local cost = 1
		
		local obs = makeObservable( priority, cost )
		
		limitedObservables[ x ] = limiter:enqueueObservable( obs, priority, cost, function()
			-- Make sure the rescheduled task is the most recently started task.
			rxTest:expectEqual( x, 5 )
		end )
		
	end
	
	-- Start the observables.
	
	for x = 1, kSimultaneousTasksLimit do
		limitedObservables[ x ] = limitedObservables[ x ]:call( runUntilDone )
		rxTest:sleep( 0.01 )
	end

	-- Now schedule a high priority task
	
	local priority = 10
	local cost = 1
	
	local obs = makeObservable( priority, cost )
	
	local highPri = limiter:enqueueObservable( obs, priority, cost ):call( runUntilDone )

	rxTest:expectOnlyCompleted( allTasksCompleted )
	
	highPri:dispose()

end

--------------------------------------------------------------------------------

return WFRxTestCase.wrapSuite( 'WFRxRateLimiterTest', WFRxRateLimiterTest )
