quantconnect--lean/Tests/Common/Scheduling/ScheduleManagerTests.cs

/*
 * QUANTCONNECT.COM - Democratizing Finance, Empowering Individuals.
 * Lean Algorithmic Trading Engine v2.0. Copyright 2014 QuantConnect Corporation.
 *
 * 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.
 *
*/

using System;
using System.Collections.Generic;
using System.Threading;
using NUnit.Framework;
using QuantConnect.Algorithm;
using QuantConnect.Lean.Engine;
using QuantConnect.Lean.Engine.DataFeeds;
using QuantConnect.Lean.Engine.RealTime;
using QuantConnect.Packets;
using QuantConnect.Scheduling;
using QuantConnect.Tests.Engine.DataFeeds;
using QuantConnect.Util.RateLimit;

namespace QuantConnect.Tests.Common.Scheduling
{
    [TestFixture, Parallelizable(ParallelScope.All)]
    public class ScheduleManagerTests
    {
        [Test]
        public void DuplicateScheduledEventsAreBothFired()
        {
            var algorithm = new QCAlgorithm();

            var handler = new BacktestingRealTimeHandler();
            var timeLimitManager = new AlgorithmTimeLimitManager(TokenBucket.Null, TimeSpan.MaxValue);
            handler.Setup(algorithm, new AlgorithmNodePacket(PacketType.BacktestNode), null, null, timeLimitManager);

            algorithm.Schedule.SetEventSchedule(handler);

            var time = new DateTime(2018, 1, 1);
            algorithm.SetDateTime(time);

            var count1 = 0;
            var count2 = 0;
            algorithm.Schedule.On(algorithm.Schedule.DateRules.EveryDay(), algorithm.Schedule.TimeRules.Every(TimeSpan.FromHours(1)), () => { count1++; });
            algorithm.Schedule.On(algorithm.Schedule.DateRules.EveryDay(), algorithm.Schedule.TimeRules.Every(TimeSpan.FromHours(1)), () => { count2++; });

            const int timeSteps = 12;

            for (var i = 0; i < timeSteps; i++)
            {
                handler.SetTime(time);
                time = time.AddHours(1);
            }

            handler.Exit();
            Assert.AreEqual(timeSteps, count1);
            Assert.AreEqual(timeSteps, count2);
        }

        [Test]
        public void TriggersWeeklyScheduledEventsEachWeekBacktesting()
        {
            var algorithm = new AlgorithmStub();

            var handler = new BacktestingRealTimeHandler();
            var time = new DateTime(2024, 02, 10);
            handler.SetTime(time);
            var timeLimitManager = new AlgorithmTimeLimitManager(TokenBucket.Null, TimeSpan.FromMinutes(20));
            handler.Setup(algorithm, new AlgorithmNodePacket(PacketType.BacktestNode), null, null, timeLimitManager);

            algorithm.Schedule.SetEventSchedule(handler);

            algorithm.SetDateTime(time);

            var spy = algorithm.AddEquity("SPY").Symbol;

            var eventTriggerTimes = new List<DateTime>();
            var scheduledEvent = algorithm.Schedule.On(algorithm.Schedule.DateRules.WeekStart(spy),
                algorithm.Schedule.TimeRules.BeforeMarketClose(spy, 3),
                () =>
                {
                    eventTriggerTimes.Add(time);
                });

            while (time.Month < 4)
            {
                handler.SetTime(time);
                time = time.AddMinutes(1);
            }

            handler.Exit();

            var expectedEventTriggerTimes = new List<DateTime>()
            {
                new DateTime(2024, 02, 12, 20, 57, 0),
                new DateTime(2024, 02, 20, 20, 57, 0),  // Monday is 19th but it's a holiday
                new DateTime(2024, 02, 26, 20, 57, 0),
                new DateTime(2024, 03, 04, 20, 57, 0),
                // Daylight saving adjustment
                new DateTime(2024, 03, 11, 19, 57, 0),
                new DateTime(2024, 03, 18, 19, 57, 0),
                new DateTime(2024, 03, 25, 19, 57, 0),
            };
            CollectionAssert.AreEqual(expectedEventTriggerTimes, eventTriggerTimes);
        }

        [Test]
        public void TriggersWeeklyScheduledEventsEachWeekLive()
        {
            var time = new DateTime(2024, 02, 10);
            SetUp(time, out var algorithm, out var handler, out var spy);
            var eventTriggerTimes = new List<DateTime>();
            var scheduledEvent = algorithm.Schedule.On(algorithm.Schedule.DateRules.WeekStart(spy),
                algorithm.Schedule.TimeRules.BeforeMarketClose(spy, 60),
                () =>
                {
                    eventTriggerTimes.Add(handler.ManualTimeProvider.GetUtcNow());
                });

            algorithm.SetFinishedWarmingUp();

            using var finished = new ManualResetEventSlim(false);

            // Schedule a task to advance time
            var timeStep = TimeSpan.FromMinutes(60);
            algorithm.Schedule.On(algorithm.Schedule.DateRules.EveryDay(),
                algorithm.Schedule.TimeRules.Every(timeStep),
                () =>
                {
                    handler.ManualTimeProvider.Advance(timeStep);
                    var now = handler.ManualTimeProvider.GetUtcNow();
                    if (now.Month >= 4)
                    {
                        finished.Set();
                    }
                });

            // Start
            handler.SetTime(time);

            // The time advance is driven by the live real-time handler firing the scheduled events above, so the
            // fast-forward from February to April is bounded by wall-clock. Give it a generous budget and assert it
            // actually reached the end, instead of asserting on a partially advanced timeline (which under load
            // dropped the final week's scheduled event and produced a misleading count mismatch).
            Assert.IsTrue(finished.Wait(TimeSpan.FromSeconds(120)),
                "Timed out waiting for the scheduled time advance to reach April");

            handler.Exit();

            var expectedEventTriggerTimes = new List<DateTime>()
            {
                new DateTime(2024, 02, 12, 20, 0, 0),
                new DateTime(2024, 02, 20, 20, 0, 0),   // Monday is 19th but it's a holiday
                new DateTime(2024, 02, 26, 20, 0, 0),
                new DateTime(2024, 03, 04, 20, 0, 0),
                // Daylight saving adjustment
                new DateTime(2024, 03, 11, 19, 0, 0),
                new DateTime(2024, 03, 18, 19, 0, 0),
                new DateTime(2024, 03, 25, 19, 0, 0),
            };
            CollectionAssert.AreEqual(expectedEventTriggerTimes, eventTriggerTimes);
        }

        [Test]
        public void DatesReturnedAreNormalized()
        {
            var time = new DateTime(2024, 02, 10);
            SetUp(time, out var algorithm, out var handler, out var spy);
            var eventTriggerTimes = new List<DateTime>();
            using var finished = new ManualResetEventSlim(false);

            // Schedule a task to advance time
            var timeStep = TimeSpan.FromMinutes(1);
            var wasCalled = false;
            Func<DateTime, DateTime, IEnumerable<DateTime>> func = (date1, date2) =>
            {
                Assert.AreEqual(DateTimeKind.Unspecified, date1.Kind);
                Assert.AreEqual(DateTimeKind.Unspecified, date2.Kind);
                wasCalled = true;
                return new List<DateTime> { date1, date2 };
            };

            algorithm.Schedule.On(new FuncDateRule("Test", func),
                algorithm.Schedule.TimeRules.Every(timeStep),
                () =>
                {
                    handler.ManualTimeProvider.Advance(timeStep);
                    var now = handler.ManualTimeProvider.GetUtcNow();
                    finished.Set();
                });

            // Start
            handler.SetTime(time);

            finished.Wait(TimeSpan.FromSeconds(15));

            handler.Exit();
            Assert.IsTrue(wasCalled);
        }

        private void SetUp(DateTime time, out QCAlgorithm algorithm, out TestableLiveTradingRealTimeHandler handler, out Symbol spy)
        {
            algorithm = new AlgorithmStub();

            handler = new TestableLiveTradingRealTimeHandler();

            handler.ManualTimeProvider.SetCurrentTime(time);
            var timeLimitManager = new AlgorithmTimeLimitManager(TokenBucket.Null, TimeSpan.FromMinutes(20));
            handler.Setup(algorithm, new LiveNodePacket(), null, null, timeLimitManager);

            algorithm.Schedule.SetEventSchedule(handler);

            algorithm.SetDateTime(time);

            spy = algorithm.AddEquity("SPY").Symbol;
        }

        private class TestableLiveTradingRealTimeHandler : LiveTradingRealTimeHandler
        {
            public ManualTimeProvider ManualTimeProvider = new ManualTimeProvider();

            protected override ITimeProvider TimeProvider => ManualTimeProvider;

            // Time is fully driven by the ManualTimeProvider in tests, so there's no need to pace the
            // scan loop to real wall-clock time. Sleeping here (Thread.Sleep granularity is ~15ms on
            // Windows) would make the loop take longer than the test's timeout for long simulated ranges,
            // causing the last scheduled events to be missed intermittently.
            protected override void WaitTillNextSecond(DateTime time)
            {
            }
        }
    }
}