diff --git a/src/WebJobs.Extensions.CosmosDB/Config/CosmosDBExtensionConfigProvider.cs b/src/WebJobs.Extensions.CosmosDB/Config/CosmosDBExtensionConfigProvider.cs index b651a851a..2b8c6134f 100644 --- a/src/WebJobs.Extensions.CosmosDB/Config/CosmosDBExtensionConfigProvider.cs +++ b/src/WebJobs.Extensions.CosmosDB/Config/CosmosDBExtensionConfigProvider.cs @@ -9,6 +9,7 @@ using Microsoft.Azure.Cosmos; using Microsoft.Azure.WebJobs.Description; using Microsoft.Azure.WebJobs.Extensions.CosmosDB.Bindings; +using Microsoft.Azure.WebJobs.Host; using Microsoft.Azure.WebJobs.Host.Bindings; using Microsoft.Azure.WebJobs.Host.Config; using Microsoft.Extensions.Logging; @@ -27,6 +28,7 @@ internal class CosmosDBExtensionConfigProvider : IExtensionConfigProvider private readonly ICosmosDBSerializerFactory _cosmosSerializerFactory; private readonly INameResolver _nameResolver; private readonly CosmosDBOptions _options; + private readonly IDrainModeManager _drainModeManager; private readonly ILoggerFactory _loggerFactory; public CosmosDBExtensionConfigProvider( @@ -34,12 +36,14 @@ public CosmosDBExtensionConfigProvider( ICosmosDBServiceFactory cosmosDBServiceFactory, ICosmosDBSerializerFactory cosmosSerializerFactory, INameResolver nameResolver, + IDrainModeManager drainModeManager, ILoggerFactory loggerFactory) { _cosmosDBServiceFactory = cosmosDBServiceFactory; _cosmosSerializerFactory = cosmosSerializerFactory; _nameResolver = nameResolver; _options = options.Value; + _drainModeManager = drainModeManager; _loggerFactory = loggerFactory; } @@ -75,7 +79,7 @@ public void Initialize(ExtensionConfigContext context) // Trigger var rule2 = context.AddBindingRule(); - rule2.BindToTrigger(new CosmosDBTriggerAttributeBindingProviderGenerator(_nameResolver, _options, this, _loggerFactory)); + rule2.BindToTrigger(new CosmosDBTriggerAttributeBindingProviderGenerator(_nameResolver, _options, this, _drainModeManager, _loggerFactory)); } internal void ValidateConnection(CosmosDBAttribute attribute, Type paramType) diff --git a/src/WebJobs.Extensions.CosmosDB/Trigger/CosmosDBTriggerAttributeBindingProvider.cs b/src/WebJobs.Extensions.CosmosDB/Trigger/CosmosDBTriggerAttributeBindingProvider.cs index 9493eb97a..3c3e6fd51 100644 --- a/src/WebJobs.Extensions.CosmosDB/Trigger/CosmosDBTriggerAttributeBindingProvider.cs +++ b/src/WebJobs.Extensions.CosmosDB/Trigger/CosmosDBTriggerAttributeBindingProvider.cs @@ -22,13 +22,18 @@ internal class CosmosDBTriggerAttributeBindingProvider private readonly CosmosDBOptions _options; private readonly ILogger _logger; private readonly CosmosDBExtensionConfigProvider _configProvider; + private readonly IDrainModeManager _drainModeManager; - public CosmosDBTriggerAttributeBindingProvider(INameResolver nameResolver, CosmosDBOptions options, - CosmosDBExtensionConfigProvider configProvider, ILoggerFactory loggerFactory) + public CosmosDBTriggerAttributeBindingProvider(INameResolver nameResolver, + CosmosDBOptions options, + CosmosDBExtensionConfigProvider configProvider, + IDrainModeManager drainModeManager, + ILoggerFactory loggerFactory) { _nameResolver = nameResolver; _options = options; _configProvider = configProvider; + _drainModeManager = drainModeManager; _logger = loggerFactory.CreateLogger(LogCategories.CreateTriggerCategory("CosmosDB")); } @@ -118,6 +123,7 @@ public async Task TryCreateAsync(TriggerBindingProviderContext monitoredContainer, leasesContainer, attribute, + _drainModeManager, _logger); } diff --git a/src/WebJobs.Extensions.CosmosDB/Trigger/CosmosDBTriggerAttributeBindingProviderGenerator.cs b/src/WebJobs.Extensions.CosmosDB/Trigger/CosmosDBTriggerAttributeBindingProviderGenerator.cs index cc81cef74..0a82854c2 100644 --- a/src/WebJobs.Extensions.CosmosDB/Trigger/CosmosDBTriggerAttributeBindingProviderGenerator.cs +++ b/src/WebJobs.Extensions.CosmosDB/Trigger/CosmosDBTriggerAttributeBindingProviderGenerator.cs @@ -4,6 +4,7 @@ using System; using System.Reflection; using System.Threading.Tasks; +using Microsoft.Azure.WebJobs.Host; using Microsoft.Azure.WebJobs.Host.Triggers; using Microsoft.Extensions.Logging; using Newtonsoft.Json.Linq; @@ -18,14 +19,19 @@ internal class CosmosDBTriggerAttributeBindingProviderGenerator : ITriggerBindin private readonly INameResolver _nameResolver; private readonly CosmosDBOptions _options; private readonly ILoggerFactory _loggerFactory; + private readonly IDrainModeManager _drainModeManager; private readonly CosmosDBExtensionConfigProvider _configProvider; - public CosmosDBTriggerAttributeBindingProviderGenerator(INameResolver nameResolver, CosmosDBOptions options, - CosmosDBExtensionConfigProvider configProvider, ILoggerFactory loggerFactory) + public CosmosDBTriggerAttributeBindingProviderGenerator(INameResolver nameResolver, + CosmosDBOptions options, + CosmosDBExtensionConfigProvider configProvider, + IDrainModeManager drainModeManager, + ILoggerFactory loggerFactory) { _nameResolver = nameResolver; _options = options; _configProvider = configProvider; + _drainModeManager = drainModeManager; _loggerFactory = loggerFactory; } @@ -57,11 +63,11 @@ public Task TryCreateAsync(TriggerBindingProviderContext contex Type genericBindingType = baseType.MakeGenericType(documentType); - Type[] typeArgs = { typeof(INameResolver), typeof(CosmosDBOptions), typeof(CosmosDBExtensionConfigProvider), typeof(ILoggerFactory) }; + Type[] typeArgs = { typeof(INameResolver), typeof(CosmosDBOptions), typeof(CosmosDBExtensionConfigProvider), typeof(IDrainModeManager), typeof(ILoggerFactory) }; ConstructorInfo constructor = genericBindingType.GetConstructor(typeArgs); - object[] constructorParameterValues = { _nameResolver, _options, _configProvider, _loggerFactory }; + object[] constructorParameterValues = { _nameResolver, _options, _configProvider, _drainModeManager, _loggerFactory }; object cosmosDBTriggerAttributeBindingProvider = constructor.Invoke(constructorParameterValues); diff --git a/src/WebJobs.Extensions.CosmosDB/Trigger/CosmosDBTriggerBinding.cs b/src/WebJobs.Extensions.CosmosDB/Trigger/CosmosDBTriggerBinding.cs index bfbdbee4c..659e30f81 100644 --- a/src/WebJobs.Extensions.CosmosDB/Trigger/CosmosDBTriggerBinding.cs +++ b/src/WebJobs.Extensions.CosmosDB/Trigger/CosmosDBTriggerBinding.cs @@ -26,6 +26,7 @@ internal class CosmosDBTriggerBinding : ITriggerBinding private readonly Container _monitoredContainer; private readonly Container _leaseContainer; private readonly CosmosDBTriggerAttribute _cosmosDBAttribute; + private readonly IDrainModeManager _drainModeManager; public CosmosDBTriggerBinding( ParameterInfo parameter, @@ -33,6 +34,7 @@ public CosmosDBTriggerBinding( Container monitoredContainer, Container leaseContainer, CosmosDBTriggerAttribute cosmosDBAttribute, + IDrainModeManager drainModeManager, ILogger logger) { _monitoredContainer = monitoredContainer; @@ -40,6 +42,7 @@ public CosmosDBTriggerBinding( _cosmosDBAttribute = cosmosDBAttribute; _parameter = parameter; _processorName = processorName; + _drainModeManager = drainModeManager; _logger = logger; } @@ -85,6 +88,7 @@ public Task CreateListenerAsync(ListenerFactoryContext context) this._monitoredContainer, this._leaseContainer, this._cosmosDBAttribute, + this._drainModeManager, this._logger)); } diff --git a/src/WebJobs.Extensions.CosmosDB/Trigger/CosmosDBTriggerListener.cs b/src/WebJobs.Extensions.CosmosDB/Trigger/CosmosDBTriggerListener.cs index 40b4ad4b8..d1db2e638 100644 --- a/src/WebJobs.Extensions.CosmosDB/Trigger/CosmosDBTriggerListener.cs +++ b/src/WebJobs.Extensions.CosmosDB/Trigger/CosmosDBTriggerListener.cs @@ -8,6 +8,7 @@ using System.Threading.Tasks; using Microsoft.Azure.Cosmos; using Microsoft.Azure.WebJobs.Extensions.CosmosDB.Trigger; +using Microsoft.Azure.WebJobs.Host; using Microsoft.Azure.WebJobs.Host.Executors; using Microsoft.Azure.WebJobs.Host.Listeners; using Microsoft.Azure.WebJobs.Host.Scale; @@ -34,6 +35,8 @@ internal class CosmosDBTriggerListener : IListener, IScaleMonitorProvider, IT private readonly string _listenerLogDetails; private readonly IScaleMonitor _cosmosDBScaleMonitor; private readonly ITargetScaler _cosmosDBTargetScaler; + private readonly CancellationTokenSource _functionExecutionCancellationTokenSource; + private readonly IDrainModeManager _drainModeManager; private ChangeFeedProcessor _host; private ChangeFeedProcessorBuilder _hostBuilder; private int _listenerStatus; @@ -45,10 +48,13 @@ public CosmosDBTriggerListener( Container monitoredContainer, Container leaseContainer, CosmosDBTriggerAttribute cosmosDBAttribute, + IDrainModeManager drainModeManager, ILogger logger) { this._logger = logger; this._executor = executor; + this._drainModeManager = drainModeManager; + this._functionExecutionCancellationTokenSource = new CancellationTokenSource(); this._processorName = processorName; this._hostName = Guid.NewGuid().ToString(); this._functionId = functionId; @@ -73,7 +79,7 @@ public void Cancel() public void Dispose() { - //Nothing to dispose + _functionExecutionCancellationTokenSource.Cancel(); } public async Task StartAsync(CancellationToken cancellationToken) @@ -118,6 +124,11 @@ public async Task StartAsync(CancellationToken cancellationToken) public async Task StopAsync(CancellationToken cancellationToken) { + if (!this._drainModeManager.IsDrainModeEnabled) + { + this._functionExecutionCancellationTokenSource.Cancel(); + } + try { if (this._host != null) @@ -211,7 +222,7 @@ internal virtual void InitializeBuilder() private async Task ProcessChangesAsync(ChangeFeedProcessorContext context, IReadOnlyCollection docs, CancellationToken cancellationToken) { this._healthMonitor.OnChangesDelivered(context); - FunctionResult result = await this._executor.TryExecuteAsync(new TriggeredFunctionData() { TriggerValue = docs }, cancellationToken); + FunctionResult result = await this._executor.TryExecuteAsync(new TriggeredFunctionData() { TriggerValue = docs }, this._functionExecutionCancellationTokenSource.Token); if (result != null // TryExecuteAsync when using RetryPolicies can return null && !result.Succeeded && result.Exception != null) @@ -220,8 +231,8 @@ private async Task ProcessChangesAsync(ChangeFeedProcessorContext context, IRead await this._healthMonitor.OnErrorAsync(context.LeaseToken, userException); } - // Prevent the change feed lease from being checkpointed if cancellation was requested - cancellationToken.ThrowIfCancellationRequested(); + // Prevent the change feed lease from being checkpointed if cancellation was requested when not in Drain mode + this._functionExecutionCancellationTokenSource.Token.ThrowIfCancellationRequested(); } public IScaleMonitor GetMonitor() diff --git a/test/WebJobs.Extensions.CosmosDB.Tests/CosmosDBConfigurationTests.cs b/test/WebJobs.Extensions.CosmosDB.Tests/CosmosDBConfigurationTests.cs index 3eaf86e7c..72e4e23a2 100644 --- a/test/WebJobs.Extensions.CosmosDB.Tests/CosmosDBConfigurationTests.cs +++ b/test/WebJobs.Extensions.CosmosDB.Tests/CosmosDBConfigurationTests.cs @@ -6,6 +6,7 @@ using System.Threading.Tasks; using Microsoft.Azure.WebJobs.Extensions.Tests.Common; using Microsoft.Azure.WebJobs.Extensions.Tests.Extensions.CosmosDB.Models; +using Microsoft.Azure.WebJobs.Host; using Microsoft.Extensions.Azure; using Microsoft.Extensions.Configuration; using Microsoft.Extensions.Logging.Abstractions; @@ -29,7 +30,7 @@ public async Task Configuration_Caches_Clients() { // Arrange var options = new CosmosDBOptions(); - var config = new CosmosDBExtensionConfigProvider(new OptionsWrapper(options), new DefaultCosmosDBServiceFactory(_baseConfig, Mock.Of()), new DefaultCosmosDBSerializerFactory(), new TestNameResolver(), NullLoggerFactory.Instance); + var config = new CosmosDBExtensionConfigProvider(new OptionsWrapper(options), new DefaultCosmosDBServiceFactory(_baseConfig, Mock.Of()), new DefaultCosmosDBSerializerFactory(), new TestNameResolver(), Mock.Of(), NullLoggerFactory.Instance); var attribute = new CosmosDBAttribute { Id = "abcdef" }; // Act @@ -104,7 +105,7 @@ private CosmosDBExtensionConfigProvider InitializeExtensionConfigProvider() var options = new CosmosDBOptions(); var factory = new DefaultCosmosDBServiceFactory(_baseConfig, Mock.Of()); var nameResolver = new TestNameResolver(); - var configProvider = new CosmosDBExtensionConfigProvider(new OptionsWrapper(options), factory, new DefaultCosmosDBSerializerFactory(), nameResolver, NullLoggerFactory.Instance); + var configProvider = new CosmosDBExtensionConfigProvider(new OptionsWrapper(options), factory, new DefaultCosmosDBSerializerFactory(), nameResolver, Mock.Of(), NullLoggerFactory.Instance); var context = TestHelpers.CreateExtensionConfigContext(nameResolver); diff --git a/test/WebJobs.Extensions.CosmosDB.Tests/CosmosDBEndToEndTests.cs b/test/WebJobs.Extensions.CosmosDB.Tests/CosmosDBEndToEndTests.cs index aebfe9b55..7cf5d8601 100644 --- a/test/WebJobs.Extensions.CosmosDB.Tests/CosmosDBEndToEndTests.cs +++ b/test/WebJobs.Extensions.CosmosDB.Tests/CosmosDBEndToEndTests.cs @@ -3,6 +3,7 @@ using System; using System.Collections.Generic; +using System.Diagnostics; using System.Linq; using System.Threading; using System.Threading.Tasks; @@ -26,6 +27,7 @@ public class CosmosDBEndToEndTests { private const string DatabaseName = "E2EDb"; private const string CollectionName = "E2ECollection"; + private const string LeaseCollectionName = "leases"; private readonly TestLoggerProvider _loggerProvider = new TestLoggerProvider(); [Fact] @@ -73,6 +75,18 @@ await TestHelpers.Await(() => .FormattedMessage; JObject loggedOptions = JObject.Parse(optionsMessage.Substring(optionsMessage.IndexOf(Environment.NewLine))); Assert.Null(loggedOptions["ConnectionMode"].Value()); + + // Clean-up leases + Container leaseContainer = client.GetContainer(DatabaseName, LeaseCollectionName); + using FeedIterator leaseIterator = leaseContainer.GetItemQueryIterator(); + while (leaseIterator.HasMoreResults) + { + FeedResponse leaseIteratorResponse = await leaseIterator.ReadNextAsync(); + foreach (JObject lease in leaseIteratorResponse) + { + await leaseContainer.DeleteItemStreamAsync(lease.Value("id"), new PartitionKey(lease.Value("id"))); + } + } } } @@ -94,6 +108,8 @@ public async Task CosmosDBEndToEndCancellation() // Start the host again and wait for the logs to show the cancelled item was reprocessed using (var host = await StartHostAsync(typeof(EndToEndCancellationTestClass))) { + var client = await InitializeDocumentClientAsync(host.Services.GetRequiredService(), DatabaseName, CollectionName); + await TestHelpers.Await(() => { var logMessages = _loggerProvider.GetAllLogMessages(); @@ -102,6 +118,18 @@ await TestHelpers.Await(() => && logMessages.Count(p => p.FormattedMessage != null && p.FormattedMessage.Contains("Saw the first document again!")) == 1 && logMessages.Count(p => p.Exception is TaskCanceledException) > 0; }); + + // Clean-up leases + Container leaseContainer = client.GetContainer(DatabaseName, LeaseCollectionName); + using FeedIterator leaseIterator = leaseContainer.GetItemQueryIterator(); + while (leaseIterator.HasMoreResults) + { + FeedResponse leaseIteratorResponse = await leaseIterator.ReadNextAsync(); + foreach (JObject lease in leaseIteratorResponse) + { + await leaseContainer.DeleteItemStreamAsync(lease.Value("id"), new PartitionKey(lease.Value("id"))); + } + } } } @@ -188,7 +216,7 @@ public static void Inputs( } public static void Trigger( - [CosmosDBTrigger(DatabaseName, CollectionName, CreateLeaseContainerIfNotExists = true)]IReadOnlyList documents, + [CosmosDBTrigger(DatabaseName, CollectionName, CreateLeaseContainerIfNotExists = true, LeaseContainerPrefix = "ciTrigger")]IReadOnlyList documents, ILogger log) { foreach (var document in documents) @@ -198,7 +226,7 @@ public static void Trigger( } public static void TriggerWithString( - [CosmosDBTrigger(DatabaseName, CollectionName, CreateLeaseContainerIfNotExists = true, LeaseContainerPrefix = "withstring")] string documents, + [CosmosDBTrigger(DatabaseName, CollectionName, CreateLeaseContainerIfNotExists = true, LeaseContainerPrefix = "ciTriggerWithString")] string documents, ILogger log) { foreach (var document in JArray.Parse(documents)) @@ -209,7 +237,7 @@ public static void TriggerWithString( [FixedDelayRetry(5, "00:00:01")] public static void TriggerWithRetry( - [CosmosDBTrigger(DatabaseName, CollectionName, CreateLeaseContainerIfNotExists = true, LeaseContainerPrefix = "retry")] IReadOnlyList documents, + [CosmosDBTrigger(DatabaseName, CollectionName, CreateLeaseContainerIfNotExists = true, LeaseContainerPrefix = "ciTriggerWithRetry")] IReadOnlyList documents, ILogger log) { foreach (var document in documents) @@ -234,7 +262,7 @@ public static async Task Trigger( DatabaseName, CollectionName, CreateLeaseContainerIfNotExists = true, - LeaseContainerPrefix = "cancellation", + LeaseContainerPrefix = "ciTriggerWithCancellation", LeaseExpirationInterval = 20 * 1000, LeaseRenewInterval = 5 * 1000, FeedPollDelay = 500, diff --git a/test/WebJobs.Extensions.CosmosDB.Tests/CosmosDBEnumerableBuilderTests.cs b/test/WebJobs.Extensions.CosmosDB.Tests/CosmosDBEnumerableBuilderTests.cs index b60462de9..19e22aaeb 100644 --- a/test/WebJobs.Extensions.CosmosDB.Tests/CosmosDBEnumerableBuilderTests.cs +++ b/test/WebJobs.Extensions.CosmosDB.Tests/CosmosDBEnumerableBuilderTests.cs @@ -10,6 +10,7 @@ using Microsoft.Azure.Cosmos; using Microsoft.Azure.WebJobs.Extensions.Tests.Common; using Microsoft.Azure.WebJobs.Extensions.Tests.Extensions.CosmosDB.Models; +using Microsoft.Azure.WebJobs.Host; using Microsoft.Extensions.Configuration; using Microsoft.Extensions.Logging.Abstractions; using Microsoft.Extensions.Options; @@ -221,7 +222,7 @@ private static CosmosDBEnumerableBuilder CreateBuilder(out Mock(new CosmosDBOptions { }); - var configProvider = new CosmosDBExtensionConfigProvider(options, mockServiceFactory.Object, new DefaultCosmosDBSerializerFactory(), new TestNameResolver(), NullLoggerFactory.Instance); + var configProvider = new CosmosDBExtensionConfigProvider(options, mockServiceFactory.Object, new DefaultCosmosDBSerializerFactory(), new TestNameResolver(), Mock.Of(), NullLoggerFactory.Instance); return new CosmosDBEnumerableBuilder(configProvider); } diff --git a/test/WebJobs.Extensions.CosmosDB.Tests/Trigger/CosmosDBListenerTests.cs b/test/WebJobs.Extensions.CosmosDB.Tests/Trigger/CosmosDBListenerTests.cs index 326db586d..97b6a2425 100644 --- a/test/WebJobs.Extensions.CosmosDB.Tests/Trigger/CosmosDBListenerTests.cs +++ b/test/WebJobs.Extensions.CosmosDB.Tests/Trigger/CosmosDBListenerTests.cs @@ -11,6 +11,7 @@ using Microsoft.Azure.Cosmos; using Microsoft.Azure.WebJobs.Extensions.CosmosDB.Trigger; using Microsoft.Azure.WebJobs.Extensions.Tests.Common; +using Microsoft.Azure.WebJobs.Host; using Microsoft.Azure.WebJobs.Host.Executors; using Microsoft.Azure.WebJobs.Host.Scale; using Microsoft.Extensions.Logging; @@ -67,7 +68,7 @@ public CosmosDBListenerTests() var attribute = new CosmosDBTriggerAttribute(DatabaseName, ContainerName); - _listener = new CosmosDBTriggerListener(_mockExecutor.Object, _functionId, ProcessorName, _monitoredContainer.Object, _leasesContainer.Object, attribute, _loggerFactory.CreateLogger>()); + _listener = new CosmosDBTriggerListener(_mockExecutor.Object, _functionId, ProcessorName, _monitoredContainer.Object, _leasesContainer.Object, attribute, Mock.Of(), _loggerFactory.CreateLogger>()); _logDetails = $"prefix='{ProcessorName}', monitoredContainer='{ContainerName}', monitoredDatabase='{DatabaseName}', " + $"leaseContainer='{ContainerName}', leaseDatabase='{DatabaseName}', functionId='{this._functionId}'"; @@ -80,7 +81,7 @@ public async Task StartAsync_Retries() var mockExecutor = new Mock(); - var listener = new MockListener(mockExecutor.Object, _functionId, ProcessorName, _monitoredContainer.Object, _leasesContainer.Object, attribute, _loggerFactory.CreateLogger>()); + var listener = new MockListener(mockExecutor.Object, _functionId, ProcessorName, _monitoredContainer.Object, _leasesContainer.Object, attribute, Mock.Of(), _loggerFactory.CreateLogger>()); // Ensure that we can call StartAsync() multiple times to retry if there is an error. for (int i = 0; i < 3; i++) @@ -117,8 +118,9 @@ public MockListener(ITriggeredFunctionExecutor executor, Container monitoredContainer, Container leaseContainer, CosmosDBTriggerAttribute cosmosDBAttribute, + IDrainModeManager drainModeManager, ILogger logger) - : base(executor, functionId, processorName, monitoredContainer, leaseContainer, cosmosDBAttribute, logger) + : base(executor, functionId, processorName, monitoredContainer, leaseContainer, cosmosDBAttribute, drainModeManager, logger) { } diff --git a/test/WebJobs.Extensions.CosmosDB.Tests/Trigger/CosmosDBTriggerAttributeBindingProviderTests.cs b/test/WebJobs.Extensions.CosmosDB.Tests/Trigger/CosmosDBTriggerAttributeBindingProviderTests.cs index 3ba1e6fab..0c31bd359 100644 --- a/test/WebJobs.Extensions.CosmosDB.Tests/Trigger/CosmosDBTriggerAttributeBindingProviderTests.cs +++ b/test/WebJobs.Extensions.CosmosDB.Tests/Trigger/CosmosDBTriggerAttributeBindingProviderTests.cs @@ -11,6 +11,7 @@ using Microsoft.Azure.WebJobs.Extensions.CosmosDB; using Microsoft.Azure.WebJobs.Extensions.CosmosDB.Tests; using Microsoft.Azure.WebJobs.Extensions.Tests.Common; +using Microsoft.Azure.WebJobs.Host; using Microsoft.Azure.WebJobs.Host.Triggers; using Microsoft.Extensions.Azure; using Microsoft.Extensions.Configuration; @@ -26,6 +27,7 @@ namespace Microsoft.Azure.WebJobs.Extensions.CosmosDBTrigger.Tests public class CosmosDBTriggerAttributeBindingProviderTests { private readonly ILoggerFactory _loggerFactory = new LoggerFactory(); + private readonly IDrainModeManager _drainModeManager = Mock.Of(); private static readonly IConfiguration _baseConfig = CosmosDBTestUtility.BuildConfiguration(new List>() { Tuple.Create(Constants.DefaultConnectionStringName, "AccountEndpoint=https://fromEnvironment;AccountKey=c29tZV9rZXk=;") @@ -83,7 +85,7 @@ public static IEnumerable ValidCosmosDBTriggerBindigsWithStartTimePara [MemberData(nameof(InvalidCosmosDBTriggerParameters))] public async Task InvalidParameters_Fail(ParameterInfo parameter) { - CosmosDBTriggerAttributeBindingProvider provider = new CosmosDBTriggerAttributeBindingProvider(new TestNameResolver(), _options, CreateExtensionConfigProvider(_options, _baseConfig), _loggerFactory); + CosmosDBTriggerAttributeBindingProvider provider = new CosmosDBTriggerAttributeBindingProvider(new TestNameResolver(), _options, CreateExtensionConfigProvider(_options, _baseConfig), _drainModeManager, _loggerFactory); InvalidOperationException ex = await Assert.ThrowsAsync(() => provider.TryCreateAsync(new TriggerBindingProviderContext(parameter, CancellationToken.None))); @@ -103,7 +105,7 @@ public async Task ValidParametersWithEnvironment_Succeed(ParameterInfo parameter }) .Build(); - CosmosDBTriggerAttributeBindingProvider provider = new CosmosDBTriggerAttributeBindingProvider(nameResolver, _options, CreateExtensionConfigProvider(_options, config), _loggerFactory); + CosmosDBTriggerAttributeBindingProvider provider = new CosmosDBTriggerAttributeBindingProvider(nameResolver, _options, CreateExtensionConfigProvider(_options, config), _drainModeManager, _loggerFactory); CosmosDBTriggerBinding binding = (CosmosDBTriggerBinding)await provider.TryCreateAsync(new TriggerBindingProviderContext(parameter, CancellationToken.None)); @@ -129,7 +131,7 @@ public async Task ValidParametersWithAppSettings_Succeed(ParameterInfo parameter }) .Build(); - CosmosDBTriggerAttributeBindingProvider provider = new CosmosDBTriggerAttributeBindingProvider(nameResolver, _options, CreateExtensionConfigProvider(_options, config), _loggerFactory); + CosmosDBTriggerAttributeBindingProvider provider = new CosmosDBTriggerAttributeBindingProvider(nameResolver, _options, CreateExtensionConfigProvider(_options, config), _drainModeManager, _loggerFactory); CosmosDBTriggerBinding binding = (CosmosDBTriggerBinding)await provider.TryCreateAsync(new TriggerBindingProviderContext(parameter, CancellationToken.None)); @@ -151,7 +153,7 @@ public async Task ValidCosmosDBTriggerBindigsWithDatabaseAndCollectionSettings_S nameResolver.Values["aDatabase"] = "myDatabase"; nameResolver.Values["aCollection"] = "myCollection"; - CosmosDBTriggerAttributeBindingProvider provider = new CosmosDBTriggerAttributeBindingProvider(nameResolver, _options, CreateExtensionConfigProvider(_options, _baseConfig), _loggerFactory); + CosmosDBTriggerAttributeBindingProvider provider = new CosmosDBTriggerAttributeBindingProvider(nameResolver, _options, CreateExtensionConfigProvider(_options, _baseConfig), _drainModeManager, _loggerFactory); CosmosDBTriggerBinding binding = (CosmosDBTriggerBinding)await provider.TryCreateAsync(new TriggerBindingProviderContext(parameter, CancellationToken.None)); @@ -180,7 +182,7 @@ public async Task ValidCosmosDBTriggerBindigsDifferentConnections_Succeed(Parame }) .Build(); - CosmosDBTriggerAttributeBindingProvider provider = new CosmosDBTriggerAttributeBindingProvider(nameResolver, _options, CreateExtensionConfigProvider(_options, config), _loggerFactory); + CosmosDBTriggerAttributeBindingProvider provider = new CosmosDBTriggerAttributeBindingProvider(nameResolver, _options, CreateExtensionConfigProvider(_options, config), _drainModeManager, _loggerFactory); CosmosDBTriggerBinding binding = (CosmosDBTriggerBinding)await provider.TryCreateAsync(new TriggerBindingProviderContext(parameter, CancellationToken.None)); @@ -204,7 +206,7 @@ public async Task ValidParametersWithEnvironment_ConnectionMode_Succeed(Paramete _options.ConnectionMode = ConnectionMode.Direct; - CosmosDBTriggerAttributeBindingProvider provider = new CosmosDBTriggerAttributeBindingProvider(nameResolver, _options, CreateExtensionConfigProvider(_options, config), _loggerFactory); + CosmosDBTriggerAttributeBindingProvider provider = new CosmosDBTriggerAttributeBindingProvider(nameResolver, _options, CreateExtensionConfigProvider(_options, config), _drainModeManager, _loggerFactory); CosmosDBTriggerBinding binding = (CosmosDBTriggerBinding)await provider.TryCreateAsync(new TriggerBindingProviderContext(parameter, CancellationToken.None)); @@ -232,7 +234,7 @@ public async Task ValidParametersWithEnvironment_UserAgent_Succeed(ParameterInfo _options.UserAgentSuffix = "randomtext"; - CosmosDBTriggerAttributeBindingProvider provider = new CosmosDBTriggerAttributeBindingProvider(nameResolver, _options, CreateExtensionConfigProvider(_options, config), _loggerFactory); + CosmosDBTriggerAttributeBindingProvider provider = new CosmosDBTriggerAttributeBindingProvider(nameResolver, _options, CreateExtensionConfigProvider(_options, config), _drainModeManager, _loggerFactory); CosmosDBTriggerBinding binding = (CosmosDBTriggerBinding)await provider.TryCreateAsync(new TriggerBindingProviderContext(parameter, CancellationToken.None)); @@ -250,7 +252,7 @@ public async Task ValidCosmosDBTriggerBindigsPreferredLocationsParameters_Succee var nameResolver = new TestNameResolver(); nameResolver.Values["regions"] = "East US, North Europe,"; - CosmosDBTriggerAttributeBindingProvider provider = new CosmosDBTriggerAttributeBindingProvider(nameResolver, _options, CreateExtensionConfigProvider(_options, _baseConfig), _loggerFactory); + CosmosDBTriggerAttributeBindingProvider provider = new CosmosDBTriggerAttributeBindingProvider(nameResolver, _options, CreateExtensionConfigProvider(_options, _baseConfig), _drainModeManager, _loggerFactory); CosmosDBTriggerBinding binding = (CosmosDBTriggerBinding)await provider.TryCreateAsync(new TriggerBindingProviderContext(parameter, CancellationToken.None)); @@ -299,7 +301,7 @@ public async Task ValidLeaseHostOptions_Succeed(ParameterInfo parameter) }) .Build(); - CosmosDBTriggerAttributeBindingProvider provider = new CosmosDBTriggerAttributeBindingProvider(nameResolver, _options, CreateExtensionConfigProvider(_options, config), _loggerFactory); + CosmosDBTriggerAttributeBindingProvider provider = new CosmosDBTriggerAttributeBindingProvider(nameResolver, _options, CreateExtensionConfigProvider(_options, config), _drainModeManager, _loggerFactory); CosmosDBTriggerBinding binding = (CosmosDBTriggerBinding)await provider.TryCreateAsync(new TriggerBindingProviderContext(parameter, CancellationToken.None)); @@ -355,7 +357,7 @@ public async Task ValidCreateIfNotExists(ParameterInfo parameter) .Setup(f => f.CreateService(It.IsAny(), It.IsAny())) .Returns(serviceMock.Object); - CosmosDBTriggerAttributeBindingProvider provider = new CosmosDBTriggerAttributeBindingProvider(new TestNameResolver(), _options, CreateExtensionConfigProvider(factoryMock.Object, _options, config), _loggerFactory); + CosmosDBTriggerAttributeBindingProvider provider = new CosmosDBTriggerAttributeBindingProvider(new TestNameResolver(), _options, CreateExtensionConfigProvider(factoryMock.Object, _options, config), _drainModeManager, _loggerFactory); CosmosDBTriggerBinding binding = (CosmosDBTriggerBinding)await provider.TryCreateAsync(new TriggerBindingProviderContext(parameter, CancellationToken.None)); @@ -413,7 +415,7 @@ public async Task ValidCreateIfNotExistsForGremlin(ParameterInfo parameter) .Setup(f => f.CreateService(It.IsAny(), It.IsAny())) .Returns(serviceMock.Object); - CosmosDBTriggerAttributeBindingProvider provider = new CosmosDBTriggerAttributeBindingProvider(new TestNameResolver(), _options, CreateExtensionConfigProvider(factoryMock.Object, _options, config), _loggerFactory); + CosmosDBTriggerAttributeBindingProvider provider = new CosmosDBTriggerAttributeBindingProvider(new TestNameResolver(), _options, CreateExtensionConfigProvider(factoryMock.Object, _options, config), _drainModeManager, _loggerFactory); CosmosDBTriggerBinding binding = (CosmosDBTriggerBinding)await provider.TryCreateAsync(new TriggerBindingProviderContext(parameter, CancellationToken.None)); @@ -438,7 +440,7 @@ public async Task ValidChangeFeedOptions_Succeed(ParameterInfo parameter) }) .Build(); - CosmosDBTriggerAttributeBindingProvider provider = new CosmosDBTriggerAttributeBindingProvider(new TestNameResolver(), _options, CreateExtensionConfigProvider(_options, config), _loggerFactory); + CosmosDBTriggerAttributeBindingProvider provider = new CosmosDBTriggerAttributeBindingProvider(new TestNameResolver(), _options, CreateExtensionConfigProvider(_options, config), _drainModeManager, _loggerFactory); CosmosDBTriggerBinding binding = (CosmosDBTriggerBinding)await provider.TryCreateAsync(new TriggerBindingProviderContext(parameter, CancellationToken.None)); @@ -464,7 +466,7 @@ public async Task ValidStartFromTime_Succeed(ParameterInfo parameter) }) .Build(); - CosmosDBTriggerAttributeBindingProvider provider = new CosmosDBTriggerAttributeBindingProvider(nameResolver, _options, CreateExtensionConfigProvider(_options, config), _loggerFactory); + CosmosDBTriggerAttributeBindingProvider provider = new CosmosDBTriggerAttributeBindingProvider(nameResolver, _options, CreateExtensionConfigProvider(_options, config), _drainModeManager, _loggerFactory); CosmosDBTriggerBinding binding = (CosmosDBTriggerBinding)await provider.TryCreateAsync(new TriggerBindingProviderContext(parameter, CancellationToken.None)); @@ -488,7 +490,7 @@ private static CosmosDBExtensionConfigProvider CreateExtensionConfigProvider(Cos private static CosmosDBExtensionConfigProvider CreateExtensionConfigProvider(ICosmosDBServiceFactory serviceFactory, CosmosDBOptions options, IConfiguration config = null) { - return new CosmosDBExtensionConfigProvider(new OptionsWrapper(options), serviceFactory, new DefaultCosmosDBSerializerFactory(), new TestNameResolver(), NullLoggerFactory.Instance); + return new CosmosDBExtensionConfigProvider(new OptionsWrapper(options), serviceFactory, new DefaultCosmosDBSerializerFactory(), new TestNameResolver(), Mock.Of(), NullLoggerFactory.Instance); } // These will use the default for ConnectionStringSetting, but override LeaseConnectionStringSetting