From ac86d94eb682356c5191c516218e4a4fd4049b60 Mon Sep 17 00:00:00 2001
From: Suhas Joshi <suhasbjoshi@hotmail.com>
Date: Wed, 17 Dec 2014 15:54:17 -0800
Subject: [PATCH] Added new methods on IdentityBuilder

---
 .../IdentityBuilder.cs                        | 10 ++++
 .../IdentityBuilderTest.cs                    | 47 ++++++++++++++++++-
 2 files changed, 56 insertions(+), 1 deletion(-)

diff --git a/src/Microsoft.AspNet.Identity/IdentityBuilder.cs b/src/Microsoft.AspNet.Identity/IdentityBuilder.cs
index 8c9d51332a..71b2011cdb 100644
--- a/src/Microsoft.AspNet.Identity/IdentityBuilder.cs
+++ b/src/Microsoft.AspNet.Identity/IdentityBuilder.cs
@@ -76,5 +76,15 @@ namespace Microsoft.AspNet.Identity
                 .AddTokenProvider(typeof(PhoneNumberTokenProvider<>).MakeGenericType(UserType))
                 .AddTokenProvider(typeof(EmailTokenProvider<>).MakeGenericType(UserType));
         }
+
+        public IdentityBuilder AddUserManager<TUserManager>() where TUserManager : class
+        {
+            return AddScoped(typeof(UserManager<>).MakeGenericType(UserType), typeof(TUserManager));
+        }
+
+        public IdentityBuilder AddRoleManager<TRoleManager>() where TRoleManager : class
+        {
+            return AddScoped(typeof(RoleManager<>).MakeGenericType(RoleType), typeof(TRoleManager));
+        }
     }
 }
\ No newline at end of file
diff --git a/test/Microsoft.AspNet.Identity.Test/IdentityBuilderTest.cs b/test/Microsoft.AspNet.Identity.Test/IdentityBuilderTest.cs
index 0bbf0f7f11..7843b95019 100644
--- a/test/Microsoft.AspNet.Identity.Test/IdentityBuilderTest.cs
+++ b/test/Microsoft.AspNet.Identity.Test/IdentityBuilderTest.cs
@@ -10,7 +10,6 @@ using System;
 using System.Threading;
 using System.Threading.Tasks;
 using System.Linq;
-using Microsoft.AspNet.Security.DataProtection;
 
 namespace Microsoft.AspNet.Identity.Test
 {
@@ -62,6 +61,28 @@ namespace Microsoft.AspNet.Identity.Test
             Assert.NotNull(thingy);
         }
 
+        [Fact]
+        public void CanOverrideUserManager()
+        {
+            var services = new ServiceCollection();
+            services.AddIdentity<TestUser, TestRole>()
+                .AddUserStore<NoopUserStore>()
+                .AddUserManager<MyUserManager>();
+            var myUserManager = services.BuildServiceProvider().GetRequiredService(typeof(UserManager<TestUser>)) as MyUserManager;
+            Assert.NotNull(myUserManager);
+        }
+
+        [Fact]
+        public void CanOverrideRoleManager()
+        {
+            var services = new ServiceCollection();
+            services.AddIdentity<TestUser, TestRole>()
+                    .AddRoleStore<NoopRoleStore>()
+                    .AddRoleManager<MyRoleManager>();
+            var myRoleManager = services.BuildServiceProvider().GetRequiredService<RoleManager<TestRole>>() as MyRoleManager;
+            Assert.NotNull(myRoleManager);
+        }
+
         [Fact]
         public void EnsureDefaultServices()
         {
@@ -203,5 +224,29 @@ namespace Microsoft.AspNet.Identity.Test
             }
         }
 
+        private class MyUserManager : UserManager<TestUser>
+        {
+            public MyUserManager(IUserStore<TestUser> store,
+            IOptions<IdentityOptions> optionsAccessor,
+            IPasswordHasher<TestUser> passwordHasher,
+            IEnumerable<IUserValidator<TestUser>> userValidators,
+            IEnumerable<IPasswordValidator<TestUser>> passwordValidators,
+            IUserNameNormalizer userNameNormalizer,
+            IEnumerable<IUserTokenProvider<TestUser>> tokenProviders,
+            IEnumerable<IIdentityMessageProvider> msgProviders) :
+                base(store, optionsAccessor, passwordHasher, userValidators, passwordValidators, userNameNormalizer, tokenProviders, msgProviders)
+            {
+
+            }
+        }
+
+        private class MyRoleManager : RoleManager<TestRole>
+        {
+            public MyRoleManager(IRoleStore<TestRole> store,
+                IEnumerable<IRoleValidator<TestRole>> roleValidators) : base(store, roleValidators)
+            {
+
+            }
+        }
     }
 }
\ No newline at end of file