quantconnect--lean/Tests/Algorithm/AlgorithmAddUniverseTests.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.Linq;
using NUnit.Framework;
using Python.Runtime;
using QuantConnect.Algorithm;
using QuantConnect.Data.Fundamental;
using QuantConnect.Data.UniverseSelection;
using QuantConnect.Tests.Engine.DataFeeds;

namespace QuantConnect.Tests.Algorithm
{
    public class AlgorithmAddUniverseTests
    {
        [TestCaseSource(nameof(ETFConstituentUniverseTestCases))]
        public void AddUniverseWithETFConstituentUniverseDefinitionTicker(string ticker, string market)
        {
            AssertConstituentUniverseDefinitionsSymbol(ticker, market, false, false, true);
        }

        [TestCaseSource(nameof(ETFConstituentUniverseTestCases))]
        public void AddUniverseWithETFConstituentUniverseDefinitionTickerPython(string ticker, string market)
        {
            AssertConstituentUniverseDefinitionsSymbol(ticker, market, false, true, true);
        }

        [TestCaseSource(nameof(ETFConstituentUniverseTestCases))]
        public void AddUniverseWithETFConstituentUniverseDefinitionSymbol(string ticker, string market)
        {
            AssertConstituentUniverseDefinitionsSymbol(ticker, market, true, false, true);
        }

        [TestCaseSource(nameof(ETFConstituentUniverseTestCases))]
        public void AddUniverseWithETFConstituentUniverseDefinitionSymbolPython(string ticker, string market)
        {
            AssertConstituentUniverseDefinitionsSymbol(ticker, market, true, true, true);
        }

        [TestCaseSource(nameof(IndexConstituentUniverseTestCases))]
        public void AddUniverseWithIndexConstituentUniverseDefinitionTicker(string ticker, string market)
        {
            AssertConstituentUniverseDefinitionsSymbol(ticker, market, false, false, false);
        }

        [TestCaseSource(nameof(IndexConstituentUniverseTestCases))]
        public void AddUniverseWithIndexConstituentUniverseDefinitionSymbol(string ticker, string market)
        {
            AssertConstituentUniverseDefinitionsSymbol(ticker, market, true, false, false);
        }

        [TestCaseSource(nameof(IndexConstituentUniverseTestCases))]
        public void AddUniverseWithIndexConstituentUniverseDefinitionTickerPython(string ticker, string market)
        {
            AssertConstituentUniverseDefinitionsSymbol(ticker, market, false, true, false);
        }

        [TestCaseSource(nameof(IndexConstituentUniverseTestCases))]
        public void AddUniverseWithIndexConstituentUniverseDefinitionSymbolPython(string ticker, string market)
        {
            AssertConstituentUniverseDefinitionsSymbol(ticker, market, true, true, false);
        }

        private void AssertConstituentUniverseDefinitionsSymbol(string ticker, string market, bool isSymbol, bool isPython, bool isEtf)
        {
            var algo = CreateAlgorithm();
            algo.SetStartDate(2021, 8, 23);
            algo.SetEndDate(2021, 8, 24);

            Universe constituentUniverse;
            var symbol = Symbol.Create(ticker, isEtf ? SecurityType.Equity : SecurityType.Index, market ?? Market.USA);

            if (isSymbol && isEtf)
            {
                constituentUniverse = isPython
                    ? algo.Universe.ETF(symbol, algo.UniverseSettings, (PyObject)null)
                    : algo.Universe.ETF(symbol, algo.UniverseSettings, CreateReturnAllFunc());
            }
            else if (isEtf)
            {
                constituentUniverse = isPython
                    ? algo.Universe.ETF(ticker, market, algo.UniverseSettings, (PyObject)null)
                    : algo.Universe.ETF(ticker, market, algo.UniverseSettings, CreateReturnAllFunc());
            }
            else if (isSymbol)
            {
                constituentUniverse = isPython
                    ? algo.Universe.Index(symbol, algo.UniverseSettings, (PyObject)null)
                    : algo.Universe.Index(symbol, algo.UniverseSettings, CreateReturnAllFunc());
            }
            else
            {
                constituentUniverse = isPython
                    ? algo.Universe.Index(ticker, market, algo.UniverseSettings, (PyObject)null)
                    : algo.Universe.Index(ticker, market, algo.UniverseSettings, CreateReturnAllFunc());
            }

            Assert.IsTrue(constituentUniverse.Configuration.Symbol.HasUnderlying);
            Assert.AreEqual(symbol, constituentUniverse.Configuration.Symbol.Underlying);

            Assert.AreEqual(symbol.SecurityType, constituentUniverse.Configuration.Symbol.SecurityType);
            Assert.IsTrue(constituentUniverse.Configuration.Symbol.ID.Symbol.StartsWithInvariant("qc-universe-"));
        }

        private static TestCaseData[] ETFConstituentUniverseTestCases()
        {
            return new[]
            {
                new TestCaseData("SPY", Market.USA),
                new TestCaseData("SPY", null),
                new TestCaseData("GDVD", Market.USA),
                new TestCaseData("GDVD", null)
            };
        }

        private static TestCaseData[] IndexConstituentUniverseTestCases()
        {
            return new[]
            {
                new TestCaseData("SPX", Market.USA),
                new TestCaseData("SPX", null),
                new TestCaseData("NDX", Market.USA),
                new TestCaseData("NDX", null)
            };
        }

        private QCAlgorithm CreateAlgorithm()
        {
            var algo = new QCAlgorithm();
            algo.SubscriptionManager.SetDataManager(new DataManagerStub(algo));
            return algo;
        }

        private Func<IEnumerable<ETFConstituentUniverse>, IEnumerable<Symbol>> CreateReturnAllFunc()
        {
            return x => x.Select(y => y.Symbol);
        }

        [TestCase(typeof(BaseDataCollection))]
        [TestCase(typeof(FundamentalUniverse))]
        [TestCase(typeof(ETFConstituentsUniverseFactory))]
        public void UniverseSymbolsSortInCreationOrder(Type dataType)
        {
            Symbol symbol1, symbol2, symbol3;
            if (dataType == typeof(ETFConstituentsUniverseFactory))
            {
                var composite = Symbol.Create("SPY", SecurityType.Equity, Market.USA);
                using var universe1 = new ETFConstituentsUniverseFactory(composite, null);
                using var universe2 = new ETFConstituentsUniverseFactory(composite, null);
                using var universe3 = new ETFConstituentsUniverseFactory(composite, null);
                symbol1 = universe1.Configuration.Symbol;
                symbol2 = universe2.Configuration.Symbol;
                symbol3 = universe3.Configuration.Symbol;
            }
            else
            {
                var instance = (BaseDataCollection)Activator.CreateInstance(dataType);
                symbol1 = instance.UniverseSymbol();
                symbol2 = instance.UniverseSymbol();
                symbol3 = instance.UniverseSymbol();
            }

            var comparer = StringComparer.OrdinalIgnoreCase;
            Assert.That(comparer.Compare(symbol1.Value, symbol2.Value) < 0);
            Assert.That(comparer.Compare(symbol2.Value, symbol3.Value) < 0);
        }
    }
}