diff --git a/src/NuGet.Services.AzureSearch/Auxiliary2AzureSearch/Auxiliary2AzureSearchCommand.cs b/src/NuGet.Services.AzureSearch/Auxiliary2AzureSearch/Auxiliary2AzureSearchCommand.cs index 0bd064d99..5ca91955e 100644 --- a/src/NuGet.Services.AzureSearch/Auxiliary2AzureSearch/Auxiliary2AzureSearchCommand.cs +++ b/src/NuGet.Services.AzureSearch/Auxiliary2AzureSearch/Auxiliary2AzureSearchCommand.cs @@ -144,8 +144,15 @@ private async Task PushIndexChangesAsync() _logger.LogInformation("Removing invalid IDs and versions from the new data."); CleanDownloadData(newData); + // Fetch the download overrides from the auxiliary file. Note that the overriden downloads are kept + // separate from downloads data as the original data will be persisted to auxiliary data, whereas the + // overriden data will be persisted to Azure Search. + _logger.LogInformation("Overriding download count data."); + var downloadOverrides = await _auxiliaryFileClient.LoadDownloadOverridesAsync(); + var overridenDownloads = newData.ApplyDownloadOverrides(downloadOverrides, _logger); + _logger.LogInformation("Detecting download count changes."); - var changes = _downloadSetComparer.Compare(oldResult.Data, newData); + var changes = _downloadSetComparer.Compare(oldResult.Data, overridenDownloads); var idBag = new ConcurrentBag(changes.Keys); _logger.LogInformation("{Count} package IDs have download count changes.", idBag.Count); diff --git a/src/NuGet.Services.AzureSearch/Auxiliary2AzureSearch/Auxiliary2AzureSearchConfiguration.cs b/src/NuGet.Services.AzureSearch/Auxiliary2AzureSearch/Auxiliary2AzureSearchConfiguration.cs index ca53293e3..7b0d2b220 100644 --- a/src/NuGet.Services.AzureSearch/Auxiliary2AzureSearch/Auxiliary2AzureSearchConfiguration.cs +++ b/src/NuGet.Services.AzureSearch/Auxiliary2AzureSearch/Auxiliary2AzureSearchConfiguration.cs @@ -11,6 +11,7 @@ public class Auxiliary2AzureSearchConfiguration : AzureSearchJobConfiguration, I public string AuxiliaryDataStorageConnectionString { get; set; } public string AuxiliaryDataStorageContainer { get; set; } public string AuxiliaryDataStorageDownloadsPath { get; set; } + public string AuxiliaryDataStorageDownloadOverridesPath { get; set; } public string AuxiliaryDataStorageExcludedPackagesPath { get; } public string AuxiliaryDataStorageVerifiedPackagesPath { get; set; } public TimeSpan MinPushPeriod { get; set; } diff --git a/src/NuGet.Services.AzureSearch/AuxiliaryFiles/AuxiliaryFileClient.cs b/src/NuGet.Services.AzureSearch/AuxiliaryFiles/AuxiliaryFileClient.cs index 20c4b2764..40c534c1a 100644 --- a/src/NuGet.Services.AzureSearch/AuxiliaryFiles/AuxiliaryFileClient.cs +++ b/src/NuGet.Services.AzureSearch/AuxiliaryFiles/AuxiliaryFileClient.cs @@ -77,6 +77,16 @@ public async Task> LoadExcludedPackagesAsync() logger: _logger)); } + public async Task> LoadDownloadOverridesAsync() + { + return await LoadAuxiliaryFileAsync( + _options.Value.AuxiliaryDataStorageDownloadOverridesPath, + loader => DownloadOverrides.Load( + fileName: null, + loader: loader, + logger: _logger)); + } + private async Task LoadAuxiliaryFileAsync( string blobName, Func loadData) where T : class diff --git a/src/NuGet.Services.AzureSearch/AuxiliaryFiles/DownloadDataExtensions.cs b/src/NuGet.Services.AzureSearch/AuxiliaryFiles/DownloadDataExtensions.cs new file mode 100644 index 000000000..144c6733a --- /dev/null +++ b/src/NuGet.Services.AzureSearch/AuxiliaryFiles/DownloadDataExtensions.cs @@ -0,0 +1,90 @@ +// Copyright (c) .NET Foundation. All rights reserved. +// Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information. + +using System; +using System.Collections.Generic; +using System.Linq; +using Microsoft.Extensions.Logging; + +namespace NuGet.Services.AzureSearch.AuxiliaryFiles +{ + public static class DownloadDataExtensions + { + public static DownloadData ApplyDownloadOverrides( + this DownloadData originalData, + IReadOnlyDictionary downloadOverrides, + ILogger logger) + { + if (originalData == null) + { + throw new ArgumentNullException(nameof(originalData)); + } + + if (downloadOverrides == null) + { + throw new ArgumentNullException(nameof(downloadOverrides)); + } + + if (logger == null) + { + throw new ArgumentNullException(nameof(logger)); + } + + // Create a copy of the original data and apply overrides as we copy. + var result = new DownloadData(); + + foreach (var downloadData in originalData) + { + var packageId = downloadData.Key; + + if (ShouldOverrideDownloads(packageId)) + { + logger.LogInformation( + "Overriding downloads of package {PackageId} from {Downloads} to {DownloadsOverride}", + packageId, + originalData.GetDownloadCount(packageId), + downloadOverrides[packageId]); + + var versions = downloadData.Value.Keys; + + result.SetDownloadCount( + packageId, + versions.First(), + downloadOverrides[packageId]); + } + else + { + foreach (var versionData in downloadData.Value) + { + result.SetDownloadCount(downloadData.Key, versionData.Key, versionData.Value); + } + } + } + + bool ShouldOverrideDownloads(string packageId) + { + if (!downloadOverrides.TryGetValue(packageId, out var downloadOverride)) + { + return false; + } + + // Apply the downloads override only if the package has fewer total downloads. + // In effect, this removes a package's manual boost once its total downloads exceed the override. + if (originalData[packageId].Total >= downloadOverride) + { + logger.LogInformation( + "Skipping download override for package {PackageId} as its downloads of {Downloads} are " + + "greater than its override of {DownloadsOverride}", + packageId, + originalData[packageId].Total, + downloadOverride); + return false; + } + + return true; + } + + return result; + } + } +} diff --git a/src/NuGet.Services.AzureSearch/AuxiliaryFiles/DownloadOverrides.cs b/src/NuGet.Services.AzureSearch/AuxiliaryFiles/DownloadOverrides.cs new file mode 100644 index 000000000..dc82206b9 --- /dev/null +++ b/src/NuGet.Services.AzureSearch/AuxiliaryFiles/DownloadOverrides.cs @@ -0,0 +1,33 @@ +using System; +using System.Collections.Generic; +using Microsoft.Extensions.Logging; +using Newtonsoft.Json; +using NuGet.Indexing; + +namespace NuGet.Services.AzureSearch.AuxiliaryFiles +{ + public static class DownloadOverrides + { + private static readonly JsonSerializer Serializer = new JsonSerializer(); + + public static IReadOnlyDictionary Load(string fileName, ILoader loader, ILogger logger) + { + try + { + using (var reader = loader.GetReader(fileName)) + { + var downloadOverrides = Serializer.Deserialize>(reader); + + return new Dictionary( + downloadOverrides, + StringComparer.OrdinalIgnoreCase); + } + } + catch (Exception ex) + { + logger.LogError(0, ex, "Unable to load download overrides {FileName} due to exception", fileName); + throw; + } + } + } +} diff --git a/src/NuGet.Services.AzureSearch/AuxiliaryFiles/IAuxiliaryDataStorageConfiguration.cs b/src/NuGet.Services.AzureSearch/AuxiliaryFiles/IAuxiliaryDataStorageConfiguration.cs index 3138af365..c394661e2 100644 --- a/src/NuGet.Services.AzureSearch/AuxiliaryFiles/IAuxiliaryDataStorageConfiguration.cs +++ b/src/NuGet.Services.AzureSearch/AuxiliaryFiles/IAuxiliaryDataStorageConfiguration.cs @@ -8,6 +8,7 @@ public interface IAuxiliaryDataStorageConfiguration string AuxiliaryDataStorageConnectionString { get; } string AuxiliaryDataStorageContainer { get; } string AuxiliaryDataStorageDownloadsPath { get; } + string AuxiliaryDataStorageDownloadOverridesPath { get; } string AuxiliaryDataStorageExcludedPackagesPath { get; } string AuxiliaryDataStorageVerifiedPackagesPath { get; } } diff --git a/src/NuGet.Services.AzureSearch/AuxiliaryFiles/IAuxiliaryFileClient.cs b/src/NuGet.Services.AzureSearch/AuxiliaryFiles/IAuxiliaryFileClient.cs index 14232f9f6..32cd90947 100644 --- a/src/NuGet.Services.AzureSearch/AuxiliaryFiles/IAuxiliaryFileClient.cs +++ b/src/NuGet.Services.AzureSearch/AuxiliaryFiles/IAuxiliaryFileClient.cs @@ -9,6 +9,7 @@ namespace NuGet.Services.AzureSearch.AuxiliaryFiles public interface IAuxiliaryFileClient { Task LoadDownloadDataAsync(); + Task> LoadDownloadOverridesAsync(); Task> LoadVerifiedPackagesAsync(); Task> LoadExcludedPackagesAsync(); } diff --git a/src/NuGet.Services.AzureSearch/Db2AzureSearch/Db2AzureSearchConfiguration.cs b/src/NuGet.Services.AzureSearch/Db2AzureSearch/Db2AzureSearchConfiguration.cs index 84c1ac32b..388d65178 100644 --- a/src/NuGet.Services.AzureSearch/Db2AzureSearch/Db2AzureSearchConfiguration.cs +++ b/src/NuGet.Services.AzureSearch/Db2AzureSearch/Db2AzureSearchConfiguration.cs @@ -13,6 +13,7 @@ public class Db2AzureSearchConfiguration : AzureSearchJobConfiguration, IAuxilia public string AuxiliaryDataStorageConnectionString { get; set; } public string AuxiliaryDataStorageContainer { get; set; } public string AuxiliaryDataStorageDownloadsPath { get; set; } + public string AuxiliaryDataStorageDownloadOverridesPath { get; set; } public string AuxiliaryDataStorageExcludedPackagesPath { get; set; } public string AuxiliaryDataStorageVerifiedPackagesPath { get; set; } } diff --git a/src/NuGet.Services.AzureSearch/Db2AzureSearch/NewPackageRegistrationProducer.cs b/src/NuGet.Services.AzureSearch/Db2AzureSearch/NewPackageRegistrationProducer.cs index 4d9d4ec70..1d93b8c69 100644 --- a/src/NuGet.Services.AzureSearch/Db2AzureSearch/NewPackageRegistrationProducer.cs +++ b/src/NuGet.Services.AzureSearch/Db2AzureSearch/NewPackageRegistrationProducer.cs @@ -54,6 +54,12 @@ public async Task ProduceWorkAsync( // numbers we don't use the gallery DB values. var downloads = await _auxiliaryFileClient.LoadDownloadDataAsync(); + // Fetch the download overrides from the auxiliary file. Note that the overriden downloads are kept + // separate from downloads data as the original data will be persisted to auxiliary data, whereas the + // overriden data will be persisted to Azure Search. + var downloadOverrides = await _auxiliaryFileClient.LoadDownloadOverridesAsync(); + var overridenDownloads = downloads.ApplyDownloadOverrides(downloadOverrides, _logger); + // Fetch the verified packages file. This is not used inside the index but is used at query-time in the // Azure Search service. We want to copy this file to the local region's storage container to improve // availability and start-up of the service. @@ -94,7 +100,7 @@ public async Task ProduceWorkAsync( allWork.Add(new NewPackageRegistration( pr.Id, - downloads.GetDownloadCount(pr.Id), + overridenDownloads.GetDownloadCount(pr.Id), pr.Owners, packages, isExcludedByDefault)); diff --git a/src/NuGet.Services.AzureSearch/NuGet.Services.AzureSearch.csproj b/src/NuGet.Services.AzureSearch/NuGet.Services.AzureSearch.csproj index aa89d1a14..0020373fb 100644 --- a/src/NuGet.Services.AzureSearch/NuGet.Services.AzureSearch.csproj +++ b/src/NuGet.Services.AzureSearch/NuGet.Services.AzureSearch.csproj @@ -52,6 +52,8 @@ + + diff --git a/tests/NuGet.Services.AzureSearch.Tests/Auxiliary2AzureSearch/Auxiliary2AzureSearchCommandFacts.cs b/tests/NuGet.Services.AzureSearch.Tests/Auxiliary2AzureSearch/Auxiliary2AzureSearchCommandFacts.cs index 4d6f6a423..f756ac814 100644 --- a/tests/NuGet.Services.AzureSearch.Tests/Auxiliary2AzureSearch/Auxiliary2AzureSearchCommandFacts.cs +++ b/tests/NuGet.Services.AzureSearch.Tests/Auxiliary2AzureSearch/Auxiliary2AzureSearchCommandFacts.cs @@ -6,7 +6,9 @@ using System.Collections.Generic; using System.Linq; using System.Threading.Tasks; +using Castle.Core.Logging; using Microsoft.Azure.Search.Models; +using Microsoft.Extensions.Logging; using Microsoft.Extensions.Options; using Moq; using NuGet.Services.AzureSearch.AuxiliaryFiles; @@ -197,6 +199,171 @@ public async Task RejectsInvalidDataAndNormalizesVersions(string propertyName) Assert.Equal(10, downloadData.GetDownloadCount("ValidId")); Assert.Contains("There were 1 invalid IDs, 2 invalid versions, and 1 non-normalized IDs.", Logger.Messages); } + + [Fact] + public async Task OverridesDownloadCounts() + { + DownloadSetComparer + .Setup(c => c.Compare(It.IsAny(), It.IsAny())) + .Returns((oldData, newData) => + { + return new SortedDictionary( + newData.ToDictionary(d => d.Key, d => d.Value.Total), + StringComparer.OrdinalIgnoreCase); + }); + + NewDownloadData.SetDownloadCount("A", "1.0.0", 12); + NewDownloadData.SetDownloadCount("A", "2.0.0", 34); + + NewDownloadData.SetDownloadCount("B", "3.0.0", 5); + NewDownloadData.SetDownloadCount("B", "4.0.0", 4); + + NewDownloadData.SetDownloadCount("C", "5.0.0", 2); + NewDownloadData.SetDownloadCount("C", "6.0.0", 3); + + DownloadOverrides["A"] = 55; + DownloadOverrides["b"] = 66; + + await Target.ExecuteAsync(); + + // Documents should have new data with overriden downloads. + SearchDocumentBuilder + .Verify( + b => b.UpdateDownloadCount("A", SearchFilters.IncludePrereleaseAndSemVer2, 55), + Times.Once); + SearchDocumentBuilder + .Verify( + b => b.UpdateDownloadCount("B", SearchFilters.IncludePrereleaseAndSemVer2, 66), + Times.Once); + SearchDocumentBuilder + .Verify( + b => b.UpdateDownloadCount("C", SearchFilters.IncludePrereleaseAndSemVer2, 5), + Times.Once); + + // Downloads auxiliary file should have new data without overriden downloads. + DownloadDataClient.Verify( + c => c.ReplaceLatestIndexedAsync( + It.Is(d => + d["A"].Total == 46 && + d["A"]["1.0.0"] == 12 && + d["A"]["2.0.0"] == 34 && + + d["B"].Total == 9 && + d["B"]["3.0.0"] == 5 && + d["B"]["4.0.0"] == 4 && + + d["C"].Total == 5 && + d["C"]["5.0.0"] == 2 && + d["C"]["6.0.0"] == 3), + It.IsAny()), + Times.Once); + } + + [Fact] + public async Task AlwaysAppliesDownloadOverrides() + { + DownloadSetComparer + .Setup(c => c.Compare(It.IsAny(), It.IsAny())) + .Returns((oldData, newData) => + { + var config = new Auxiliary2AzureSearchConfiguration(); + var telemetry = Mock.Of(); + var logger = Mock.Of>(); + var options = new Mock>(); + + options.Setup(o => o.Value).Returns(config); + + return new DownloadSetComparer(telemetry, options.Object, logger) + .Compare(oldData, newData); + }); + + // Download override should be applied even if the package's downloads haven't changed. + OldDownloadData.SetDownloadCount("A", "1.0.0", 1); + NewDownloadData.SetDownloadCount("A", "1.0.0", 1); + DownloadOverrides["A"] = 2; + + await Target.ExecuteAsync(); + + // Documents should have new data with overriden downloads. + SearchDocumentBuilder + .Verify( + b => b.UpdateDownloadCount("A", SearchFilters.IncludePrereleaseAndSemVer2, 2), + Times.Once); + + // Downloads auxiliary file should have new data without overriden downloads. + DownloadDataClient.Verify( + c => c.ReplaceLatestIndexedAsync( + It.Is(d => + d["A"].Total == 1 && + d["A"]["1.0.0"] == 1), + It.IsAny()), + Times.Once); + } + + [Fact] + public async Task DoesNotOverrideIfDownloadsGreaterOrPackageHasNoDownloads() + { + DownloadSetComparer + .Setup(c => c.Compare(It.IsAny(), It.IsAny())) + .Returns((oldData, newData) => + { + return new SortedDictionary( + newData.ToDictionary(d => d.Key, d => d.Value.Total), + StringComparer.OrdinalIgnoreCase); + }); + + NewDownloadData.SetDownloadCount("A", "1.0.0", 100); + NewDownloadData.SetDownloadCount("A", "2.0.0", 200); + + NewDownloadData.SetDownloadCount("B", "3.0.0", 5); + NewDownloadData.SetDownloadCount("B", "4.0.0", 4); + + NewDownloadData.SetDownloadCount("C", "5.0.0", 0); + + DownloadOverrides["A"] = 55; + DownloadOverrides["C"] = 66; + DownloadOverrides["D"] = 77; + + await Target.ExecuteAsync(); + + // Documents should have new data with overriden downloads. + SearchDocumentBuilder + .Verify( + b => b.UpdateDownloadCount("A", SearchFilters.IncludePrereleaseAndSemVer2, 300), + Times.Once); + SearchDocumentBuilder + .Verify( + b => b.UpdateDownloadCount("B", SearchFilters.IncludePrereleaseAndSemVer2, 9), + Times.Once); + SearchDocumentBuilder + .Verify( + b => b.UpdateDownloadCount("B", SearchFilters.IncludePrereleaseAndSemVer2, 9), + Times.Once); + SearchDocumentBuilder + .Verify( + b => b.UpdateDownloadCount("C", It.IsAny(), It.IsAny()), + Times.Never); + SearchDocumentBuilder + .Verify( + b => b.UpdateDownloadCount("D", It.IsAny(), It.IsAny()), + Times.Never); + + // Downloads auxiliary file should have new data without overriden downloads. + DownloadDataClient.Verify( + c => c.ReplaceLatestIndexedAsync( + It.Is(d => + d.Keys.Count() == 2 && + + d["A"].Total == 300 && + d["A"]["1.0.0"] == 100 && + d["A"]["2.0.0"] == 200 && + + d["B"].Total == 9 && + d["B"]["3.0.0"] == 5 && + d["B"]["4.0.0"] == 4), + It.IsAny()), + Times.Once); + } } public abstract class Facts @@ -231,6 +398,10 @@ public Facts(ITestOutputHelper output) .ReturnsAsync(() => OldDownloadResult); NewDownloadData = new DownloadData(); AuxiliaryFileClient.Setup(x => x.LoadDownloadDataAsync()).ReturnsAsync(() => NewDownloadData); + DownloadOverrides = new Dictionary(StringComparer.OrdinalIgnoreCase); + AuxiliaryFileClient + .Setup(x => x.LoadDownloadOverridesAsync()) + .ReturnsAsync(() => DownloadOverrides); OldVerifiedPackagesData = new HashSet(); OldVerifiedPackagesResult = Data.GetAuxiliaryFileResult(OldVerifiedPackagesData, "verified-packages-etag"); @@ -310,6 +481,7 @@ public Facts(ITestOutputHelper output) public DownloadData OldDownloadData { get; } public AuxiliaryFileResult OldDownloadResult { get; } public DownloadData NewDownloadData { get; } + public Dictionary DownloadOverrides { get; } public SortedDictionary Changes { get; } public Auxiliary2AzureSearchCommand Target { get; } public IndexActions IndexActions { get; set; } diff --git a/tests/NuGet.Services.AzureSearch.Tests/AuxiliaryFiles/AuxiliaryFileClientFacts.cs b/tests/NuGet.Services.AzureSearch.Tests/AuxiliaryFiles/AuxiliaryFileClientFacts.cs index ee27d9d31..f930a581c 100644 --- a/tests/NuGet.Services.AzureSearch.Tests/AuxiliaryFiles/AuxiliaryFileClientFacts.cs +++ b/tests/NuGet.Services.AzureSearch.Tests/AuxiliaryFiles/AuxiliaryFileClientFacts.cs @@ -147,6 +147,40 @@ public async Task ThrowsStorageExceptionWhenNotFound() } } + public class LoadDownloadOverridesAsync : BaseFacts + { + public LoadDownloadOverridesAsync(ITestOutputHelper output) : base(output) + { + } + + [Fact] + public async Task ReadsContent() + { + var json = @" +{ + ""A"": 10, + ""B"": 1 +} +"; + _blob + .Setup(x => x.OpenReadAsync(It.IsAny())) + .ReturnsAsync(() => new MemoryStream(Encoding.UTF8.GetBytes(json))); + + var actual = await _target.LoadDownloadOverridesAsync(); + + Assert.NotNull(actual); + Assert.Equal(10, actual["A"]); + Assert.Equal(1, actual["B"]); + Assert.Equal(1, actual["b"]); + _blobClient.Verify(x => x.GetContainerReference("my-container"), Times.Once); + _blobClient.Verify(x => x.GetContainerReference(It.IsAny()), Times.Once); + _container.Verify(x => x.GetBlobReference("my-download-overrides.json"), Times.Once); + _container.Verify(x => x.GetBlobReference(It.IsAny()), Times.Once); + _blob.Verify(x => x.OpenReadAsync(It.Is(a => a.IfMatchETag == null && a.IfNoneMatchETag == null)), Times.Once); + _blob.Verify(x => x.OpenReadAsync(It.IsAny()), Times.Once); + } + } + public abstract class BaseFacts { protected readonly Mock _blobClient; @@ -172,6 +206,7 @@ public BaseFacts(ITestOutputHelper output) _config.AuxiliaryDataStorageContainer = "my-container"; _config.AuxiliaryDataStorageDownloadsPath = "my-downloads.json"; + _config.AuxiliaryDataStorageDownloadOverridesPath = "my-download-overrides.json"; _config.AuxiliaryDataStorageVerifiedPackagesPath = "my-verified-packages.json"; _config.AuxiliaryDataStorageExcludedPackagesPath = "my-excluded-packages.json"; _options.Setup(x => x.Value).Returns(() => _config); diff --git a/tests/NuGet.Services.AzureSearch.Tests/Db2AzureSearch/NewPackageRegistrationProducerFacts.cs b/tests/NuGet.Services.AzureSearch.Tests/Db2AzureSearch/NewPackageRegistrationProducerFacts.cs index bc01db657..87ca0eb77 100644 --- a/tests/NuGet.Services.AzureSearch.Tests/Db2AzureSearch/NewPackageRegistrationProducerFacts.cs +++ b/tests/NuGet.Services.AzureSearch.Tests/Db2AzureSearch/NewPackageRegistrationProducerFacts.cs @@ -37,6 +37,7 @@ public class ProduceWorkAsync private readonly NewPackageRegistrationProducer _target; private readonly Mock _auxiliaryFileClient; private readonly DownloadData _downloads; + private readonly Dictionary _downloadOverrides; private readonly HashSet _verifiedPackages; private HashSet _excludedPackages; @@ -64,6 +65,10 @@ public ProduceWorkAsync(ITestOutputHelper output) _auxiliaryFileClient .Setup(x => x.LoadDownloadDataAsync()) .ReturnsAsync(() => _downloads); + _downloadOverrides = new Dictionary(StringComparer.OrdinalIgnoreCase); + _auxiliaryFileClient + .Setup(x => x.LoadDownloadOverridesAsync()) + .ReturnsAsync(() => _downloadOverrides); _verifiedPackages = new HashSet(); _auxiliaryFileClient .Setup(x => x.LoadVerifiedPackagesAsync()) @@ -364,6 +369,154 @@ public async Task ThrowsWhenExcludedPackagesIsMissing() await Assert.ThrowsAsync(async () => await _target.ProduceWorkAsync(_work, _token)); } + [Fact] + public async Task OverridesDownloadCounts() + { + _packageRegistrations.Add(new PackageRegistration + { + Key = 1, + Id = "A", + Packages = new[] + { + new Package { Version = "1.0.0" }, + new Package { Version = "2.0.0" }, + }, + }); + _downloads.SetDownloadCount("A", "1.0.0", 12); + _downloads.SetDownloadCount("A", "2.0.0", 23); + _packageRegistrations.Add(new PackageRegistration + { + Key = 2, + Id = "B", + Packages = new[] + { + new Package { Version = "3.0.0" }, + new Package { Version = "4.0.0" }, + }, + }); + _downloads.SetDownloadCount("B", "3.0.0", 5); + _downloads.SetDownloadCount("B", "4.0.0", 4); + _packageRegistrations.Add(new PackageRegistration + { + Key = 3, + Id = "C", + Packages = new[] + { + new Package { Version = "5.0.0" }, + new Package { Version = "6.0.0" }, + }, + }); + _downloads.SetDownloadCount("C", "5.0.0", 2); + _downloads.SetDownloadCount("C", "6.0.0", 3); + + InitializePackagesFromPackageRegistrations(); + + _downloadOverrides["A"] = 55; + _downloadOverrides["b"] = 66; + + var result = await _target.ProduceWorkAsync(_work, _token); + + // Documents should have overriden downloads. + var work = _work.Reverse().ToList(); + Assert.Equal(3, work.Count); + + Assert.Equal("A", work[0].PackageId); + Assert.Equal("1.0.0", work[0].Packages[0].Version); + Assert.Equal("2.0.0", work[0].Packages[1].Version); + Assert.Equal(55, work[0].TotalDownloadCount); + + Assert.Equal("B", work[1].PackageId); + Assert.Equal("3.0.0", work[1].Packages[0].Version); + Assert.Equal("4.0.0", work[1].Packages[1].Version); + Assert.Equal(66, work[1].TotalDownloadCount); + + Assert.Equal("C", work[2].PackageId); + Assert.Equal("5.0.0", work[2].Packages[0].Version); + Assert.Equal("6.0.0", work[2].Packages[1].Version); + Assert.Equal(5, work[2].TotalDownloadCount); + + // Downloads auxiliary file should have original downloads. + Assert.Equal(12, result.Downloads["A"]["1.0.0"]); + Assert.Equal(23, result.Downloads["A"]["2.0.0"]); + Assert.Equal(5, result.Downloads["B"]["3.0.0"]); + Assert.Equal(4, result.Downloads["B"]["4.0.0"]); + Assert.Equal(2, result.Downloads["C"]["5.0.0"]); + Assert.Equal(3, result.Downloads["C"]["6.0.0"]); + } + + [Fact] + public async Task DoesNotOverrideIfDownloadsGreaterOrPackageHasNoDownloads() + { + _packageRegistrations.Add(new PackageRegistration + { + Key = 1, + Id = "A", + Packages = new[] + { + new Package { Version = "1.0.0" }, + new Package { Version = "2.0.0" }, + }, + }); + _downloads.SetDownloadCount("A", "1.0.0", 100); + _downloads.SetDownloadCount("A", "2.0.0", 200); + _packageRegistrations.Add(new PackageRegistration + { + Key = 2, + Id = "B", + Packages = new[] + { + new Package { Version = "3.0.0" }, + new Package { Version = "4.0.0" }, + }, + }); + _downloads.SetDownloadCount("B", "3.0.0", 5); + _downloads.SetDownloadCount("B", "4.0.0", 4); + _packageRegistrations.Add(new PackageRegistration + { + Key = 3, + Id = "C", + Packages = new[] + { + new Package { Version = "5.0.0" }, + }, + }); + _downloads.SetDownloadCount("C", "5.0.0", 0); + + InitializePackagesFromPackageRegistrations(); + + _downloadOverrides["A"] = 55; + _downloadOverrides["C"] = 66; + _downloadOverrides["D"] = 77; + + var result = await _target.ProduceWorkAsync(_work, _token); + + // Documents should have overriden downloads. + var work = _work.Reverse().ToList(); + Assert.Equal(3, work.Count); + + Assert.Equal("A", work[0].PackageId); + Assert.Equal("1.0.0", work[0].Packages[0].Version); + Assert.Equal("2.0.0", work[0].Packages[1].Version); + Assert.Equal(300, work[0].TotalDownloadCount); + + Assert.Equal("B", work[1].PackageId); + Assert.Equal("3.0.0", work[1].Packages[0].Version); + Assert.Equal("4.0.0", work[1].Packages[1].Version); + Assert.Equal(9, work[1].TotalDownloadCount); + + Assert.Equal("C", work[2].PackageId); + Assert.Equal("5.0.0", work[2].Packages[0].Version); + Assert.Equal(0, work[2].TotalDownloadCount); + + // Downloads auxiliary file should have original downloads. + Assert.Equal(100, result.Downloads["A"]["1.0.0"]); + Assert.Equal(200, result.Downloads["A"]["2.0.0"]); + Assert.Equal(5, result.Downloads["B"]["3.0.0"]); + Assert.Equal(4, result.Downloads["B"]["4.0.0"]); + Assert.DoesNotContain("C", result.Downloads.Keys); + Assert.DoesNotContain("D", result.Downloads.Keys); + } + private void InitializePackagesFromPackageRegistrations() { foreach (var pr in _packageRegistrations)