Make sure that that SNS & SQS Handler are using scoped. (#15)

So after usage scoped objects are disposed. Co-authored-by: Wouter Crooy <[email protected]>
Kralizek · May 4, 2020 · 662bb12 · 662bb12
1 parent c576de1
commit 662bb12
Show file tree

Hide file tree

Showing 6 changed files with 198 additions and 8 deletions.
diff --git a/src/Kralizek.Lambda.Template.Sns/SnsEventHandler.cs b/src/Kralizek.Lambda.Template.Sns/SnsEventHandler.cs
@@ -25,12 +25,12 @@ public async Task HandleAsync(SNSEvent input, ILambdaContext context)
         {
             foreach (var record in input.Records)
             {
-                using (_serviceProvider.CreateScope())
+                using (var scope = _serviceProvider.CreateScope())
                 {
                     var message = record.Sns.Message;
                     var notification = JsonSerializer.Deserialize<TNotification>(message);
-
-                    var handler = _serviceProvider.GetService<INotificationHandler<TNotification>>();
+                    
+                    var handler = scope.ServiceProvider.GetService<INotificationHandler<TNotification>>();
 
                     if (handler == null)
                     {

diff --git a/src/Kralizek.Lambda.Template.Sqs/SqsEventHandler.cs b/src/Kralizek.Lambda.Template.Sqs/SqsEventHandler.cs
@@ -23,12 +23,12 @@ public async Task HandleAsync(SQSEvent input, ILambdaContext context)
         {
             foreach (var record in input.Records)
             {
-                using (_serviceProvider.CreateScope())
+                using (var scope = _serviceProvider.CreateScope())
                 {
                     var sqsMessage = record.Body;
                     var message = JsonSerializer.Deserialize<TMessage>(sqsMessage);
 
-                    var handler = _serviceProvider.GetService<IMessageHandler<TMessage>>();
+                    var handler = scope.ServiceProvider.GetService<IMessageHandler<TMessage>>();
 
                     if (handler == null)
                     {

diff --git a/tests/Tests.Lambda.Template/Sns/SnsEventHandlerDisposalTests.cs b/tests/Tests.Lambda.Template/Sns/SnsEventHandlerDisposalTests.cs
@@ -0,0 +1,90 @@
+using System;
+using System.Collections.Generic;
+using System.Threading.Tasks;
+using Amazon.Lambda.Core;
+using Amazon.Lambda.SNSEvents;
+using Amazon.Lambda.SQSEvents;
+using Amazon.Lambda.TestUtilities;
+using Kralizek.Lambda;
+using Microsoft.Extensions.DependencyInjection;
+using Microsoft.Extensions.Logging.Abstractions;
+using NUnit.Framework;
+
+namespace Tests.Lambda.Sns
+{
+    public class SnsEventHandlerDisposalTests
+    {
+        [Test]
+        public async Task EventHandler_Should_Use_Scoped_Object_In_ForEach_Loop()
+        {
+            var snsEvent = new SNSEvent
+            {
+                Records = new List<SNSEvent.SNSRecord>
+                {
+                    new SNSEvent.SNSRecord
+                    {
+                        Sns = new SNSEvent.SNSMessage
+                        {
+                            Message = "{}"
+                        }
+                    },
+                    new SNSEvent.SNSRecord
+                    {
+                        Sns = new SNSEvent.SNSMessage
+                        {
+                            Message = "{}"
+                        }
+                    }
+                }
+            };
+
+            var dependency = new DisposableDependency();
+
+            var services = new ServiceCollection();
+
+            services.AddScoped(_ => dependency);
+
+            var tcs = new TaskCompletionSource<TestNotification>();
+            services.AddTransient<IEventHandler<SQSEvent>, SqsEventHandler<TestNotification>>();
+
+            services.AddTransient<INotificationHandler<TestNotification>,
+                TestNotificationScopedHandler>(provider =>
+                new TestNotificationScopedHandler(provider.GetRequiredService<DisposableDependency>(), tcs));
+
+            var sp = services.BuildServiceProvider();
+            var snsEventHandler = new SnsEventHandler<TestNotification>(sp, new NullLoggerFactory());
+
+            var task = snsEventHandler.HandleAsync(snsEvent, new TestLambdaContext());
+
+            Assert.That(dependency.Disposed, Is.False, "Dependency should not be disposed");
+            Assert.That(task.IsCompleted, Is.False, "The task should not be completed");
+
+            tcs.SetResult(new TestNotification());
+
+            await task;
+
+            Assert.That(dependency.Disposed, Is.True, "Dependency should be disposed");
+            Assert.That(task.IsCompleted, Is.True, "The task should be completed");
+        }
+
+        private class DisposableDependency : IDisposable
+        {
+            public bool Disposed { get; private set; }
+            public void Dispose() => Disposed = true;
+        }
+
+        private class TestNotificationScopedHandler: INotificationHandler<TestNotification>
+        {
+            private readonly DisposableDependency _dependency;
+            private readonly TaskCompletionSource<TestNotification> _tcs;
+
+            public TestNotificationScopedHandler(DisposableDependency dependency, TaskCompletionSource<TestNotification> tcs)
+            {
+                _dependency = dependency;
+                _tcs = tcs;
+            }
+
+            public Task HandleAsync(TestNotification message, ILambdaContext context) => _tcs.Task;
+        }
+    }
+}
diff --git a/tests/Tests.Lambda.Template/Sns/SnsEventHandlerTests.cs b/tests/Tests.Lambda.Template/Sns/SnsEventHandlerTests.cs
@@ -20,23 +20,30 @@ public class SnsEventHandlerTests
         private Mock<IServiceScopeFactory> mockServiceScopeFactory;
         private Mock<IServiceProvider> mockServiceProvider;
         private Mock<ILoggerFactory> mockLoggerFactory;
+        private Mock<IServiceScope> mockServiceScope;
+
 
         [SetUp]
         public void Initialize()
         {
             mockNotificationHandler = new Mock<INotificationHandler<TestNotification>>();
             mockNotificationHandler.Setup(p => p.HandleAsync(It.IsAny<TestNotification>(), It.IsAny<ILambdaContext>()))
                                     .Returns(Task.CompletedTask);
-
+
+            mockServiceScope = new Mock<IServiceScope>();
+
             mockServiceScopeFactory = new Mock<IServiceScopeFactory>();
-            mockServiceScopeFactory.Setup(p => p.CreateScope()).Returns(Mock.Of<IServiceScope>());
+            mockServiceScopeFactory.Setup(p => p.CreateScope()).Returns(mockServiceScope.Object);
 
             mockServiceProvider = new Mock<IServiceProvider>();
             mockServiceProvider.Setup(p => p.GetService(typeof(INotificationHandler<TestNotification>)))
                                     .Returns(mockNotificationHandler.Object);
             mockServiceProvider.Setup(p => p.GetService(typeof(IServiceScopeFactory)))
                                     .Returns(mockServiceScopeFactory.Object);
 
+            mockServiceScope.Setup(p => p.ServiceProvider).Returns(mockServiceProvider.Object);
+
+
             mockLoggerFactory = new Mock<ILoggerFactory>();
             mockLoggerFactory.Setup(p => p.CreateLogger(It.IsAny<string>()))
                                 .Returns(Mock.Of<ILogger>());
@@ -174,6 +181,7 @@ public void HandleAsync_throws_InvalidOperation_if_NotificationHandler_is_not_re
 
             mockServiceProvider = new Mock<IServiceProvider>();
             mockServiceProvider.Setup(p => p.GetService(typeof(IServiceScopeFactory))).Returns(mockServiceScopeFactory.Object);
+            mockServiceScope.Setup(p => p.ServiceProvider).Returns(mockServiceProvider.Object);
 
             var sut = CreateSystemUnderTest();
 

diff --git a/tests/Tests.Lambda.Template/Sqs/SqsEventHandlerDisposalTests.cs b/tests/Tests.Lambda.Template/Sqs/SqsEventHandlerDisposalTests.cs
@@ -0,0 +1,83 @@
+using System;
+using System.Collections.Generic;
+using System.Threading.Tasks;
+using Amazon.Lambda.Core;
+using Amazon.Lambda.SQSEvents;
+using Amazon.Lambda.TestUtilities;
+using Kralizek.Lambda;
+using Microsoft.Extensions.DependencyInjection;
+using Microsoft.Extensions.Logging.Abstractions;
+using NUnit.Framework;
+
+namespace Tests.Lambda.Sqs
+{
+    public class SqsEventHandlerDisposalTests
+    {
+        [Test]
+        public async Task EventHandler_Should_Use_Scoped_Object_In_ForEach_Loop()
+        {
+            var sqsEvent = new SQSEvent
+            {
+                Records = new List<SQSEvent.SQSMessage>
+                {
+                    new SQSEvent.SQSMessage
+                    {
+                        Body = "{}"
+                    },
+                    new SQSEvent.SQSMessage
+                    {
+                        Body = "{}"
+                    },
+                }
+            };
+
+            var dependency = new DisposableDependency();
+
+            var services = new ServiceCollection();
+
+            services.AddScoped(_ => dependency);
+
+            var tcs = new TaskCompletionSource<TestMessage>();
+            services.AddTransient<IEventHandler<SQSEvent>, SqsEventHandler<TestMessage>>();
+
+            services.AddTransient<IMessageHandler<TestMessage>,
+                TestMessageScopedHandler>(provider =>
+                new TestMessageScopedHandler(provider.GetRequiredService<DisposableDependency>(), tcs));
+
+            var sp = services.BuildServiceProvider();
+            var sqsEventHandler = new SqsEventHandler<TestMessage>(sp, new NullLoggerFactory());
+
+            var task = sqsEventHandler.HandleAsync(sqsEvent, new TestLambdaContext());
+
+            Assert.That(dependency.Disposed, Is.False, "Dependency should not be disposed");
+            Assert.That(task.IsCompleted, Is.False, "The task should not be completed");
+
+            tcs.SetResult(new TestMessage());
+            await task;
+            Assert.That(dependency.Disposed, Is.True, "Dependency should be disposed");
+            Assert.That(task.IsCompleted, Is.True, "The task should be completed");
+
+
+        }
+
+        private class DisposableDependency : IDisposable
+        {
+            public bool Disposed { get; private set; }
+            public void Dispose() => Disposed = true;
+        }
+
+        private class TestMessageScopedHandler : IMessageHandler<TestMessage>
+        {
+            private readonly DisposableDependency _dependency;
+            private readonly TaskCompletionSource<TestMessage> _tcs;
+
+            public TestMessageScopedHandler(DisposableDependency dependency, TaskCompletionSource<TestMessage> tcs)
+            {
+                _dependency = dependency;
+                _tcs = tcs;
+            }
+
+            public Task HandleAsync(TestMessage message, ILambdaContext context) => _tcs.Task;
+        }
+    }
+}
diff --git a/tests/Tests.Lambda.Template/Sqs/SqsEventHandlerTests.cs b/tests/Tests.Lambda.Template/Sqs/SqsEventHandlerTests.cs
@@ -7,6 +7,7 @@
 using Kralizek.Lambda;
 using Microsoft.Extensions.DependencyInjection;
 using Microsoft.Extensions.Logging;
+using Microsoft.Extensions.Logging.Abstractions;
 using Moq;
 using NUnit.Framework;
 
@@ -18,6 +19,7 @@ public class SqsEventHandlerTests
         private Mock<IServiceScopeFactory> mockServiceScopeFactory;
         private Mock<IServiceProvider> mockServiceProvider;
         private Mock<ILoggerFactory> mockLoggerFactory;
+        private Mock<IServiceScope> mockServiceScope;
 
 
         [SetUp]
@@ -26,15 +28,20 @@ public void Initialize()
             mockMessageHandler = new Mock<IMessageHandler<TestMessage>>();
             mockMessageHandler.Setup(p => p.HandleAsync(It.IsAny<TestMessage>(), It.IsAny<ILambdaContext>())).Returns(Task.CompletedTask);
 
+            mockServiceScope = new Mock<IServiceScope>();
+
             mockServiceScopeFactory = new Mock<IServiceScopeFactory>();
-            mockServiceScopeFactory.Setup(p => p.CreateScope()).Returns(Mock.Of<IServiceScope>());
+
+            mockServiceScopeFactory.Setup(p => p.CreateScope()).Returns(mockServiceScope.Object);
 
             mockServiceProvider = new Mock<IServiceProvider>();
             mockServiceProvider.Setup(p => p.GetService(typeof(IMessageHandler<TestMessage>)))
                                .Returns(mockMessageHandler.Object);
             mockServiceProvider.Setup(p => p.GetService(typeof(IServiceScopeFactory)))
                                .Returns(mockServiceScopeFactory.Object);
 
+            mockServiceScope.Setup(p => p.ServiceProvider).Returns(mockServiceProvider.Object);
+
             mockLoggerFactory = new Mock<ILoggerFactory>();
             mockLoggerFactory.Setup(p => p.CreateLogger(It.IsAny<string>()))
                              .Returns(Mock.Of<ILogger>());
@@ -148,6 +155,8 @@ public void HandleAsync_throws_InvalidOperation_if_NotificationHandler_is_not_re
 
             mockServiceProvider = new Mock<IServiceProvider>();
             mockServiceProvider.Setup(p => p.GetService(typeof(IServiceScopeFactory))).Returns(mockServiceScopeFactory.Object);
+
+            mockServiceScope.Setup(p => p.ServiceProvider).Returns(mockServiceProvider.Object);
 
             var sut = CreateSystemUnderTest();