From 5418d817d002f8946168b924ea535ec919565ec2 Mon Sep 17 00:00:00 2001
From: Todd Martin <tmartin@bitwarden.com>
Date: Fri, 13 Jun 2025 19:16:00 -0400
Subject: [PATCH] Fixed tests.

---
 .../Services/TwoFactorEmailServiceTests.cs    | 64 +++++++++++--------
 1 file changed, 38 insertions(+), 26 deletions(-)

diff --git a/test/Core.Test/Auth/Services/TwoFactorEmailServiceTests.cs b/test/Core.Test/Auth/Services/TwoFactorEmailServiceTests.cs
index efd5160a9c..0eb72900ae 100644
--- a/test/Core.Test/Auth/Services/TwoFactorEmailServiceTests.cs
+++ b/test/Core.Test/Auth/Services/TwoFactorEmailServiceTests.cs
@@ -18,14 +18,17 @@ namespace Bit.Core.Test.Auth.Services;
 [SutProviderCustomize]
 public class TwoFactorEmailServiceTests
 {
-
     [Theory, BitAutoData]
     public async Task SendTwoFactorEmailAsync_Success(SutProvider<TwoFactorEmailService> sutProvider, User user)
     {
         var email = user.Email.ToLowerInvariant();
         var token = "thisisatokentocompare";
         var IpAddress = "1.1.1.1";
-        var deviceType = "Android";
+        var deviceType = DeviceType.Android;
+
+        var context = sutProvider.GetDependency<ICurrentContext>();
+        context.DeviceType = deviceType;
+        context.IpAddress = IpAddress;
 
         var userTwoFactorTokenProvider = Substitute.For<IUserTwoFactorTokenProvider<User>>();
         userTwoFactorTokenProvider
@@ -35,10 +38,6 @@ public class TwoFactorEmailServiceTests
             .GenerateAsync("TwoFactor", Arg.Any<UserManager<User>>(), user)
             .Returns(Task.FromResult(token));
 
-        var context = Substitute.For<ICurrentContext>();
-        context.DeviceType = DeviceType.Android;
-        context.IpAddress = IpAddress;
-
         var userManager = sutProvider.GetDependency<UserManager<User>>();
         userManager.RegisterTokenProvider(CoreHelpers.CustomProviderName(TwoFactorProviderType.Email), userTwoFactorTokenProvider);
 
@@ -54,7 +53,7 @@ public class TwoFactorEmailServiceTests
 
         await sutProvider.GetDependency<IMailService>()
             .Received(1)
-            .SendTwoFactorEmailAsync(email, user.Email, token, IpAddress, deviceType,
+            .SendTwoFactorEmailAsync(email, user.Email, token, IpAddress, deviceType.ToString(),
                 TwoFactorEmailPurpose.Login);
     }
 
@@ -64,7 +63,11 @@ public class TwoFactorEmailServiceTests
         var email = user.Email.ToLowerInvariant();
         var token = "thisisatokentocompare";
         var IpAddress = "1.1.1.1";
-        var deviceType = "Android";
+        var deviceType = DeviceType.Android;
+
+        var context = sutProvider.GetDependency<ICurrentContext>();
+        context.DeviceType = deviceType;
+        context.IpAddress = IpAddress;
 
         var userTwoFactorTokenProvider = Substitute.For<IUserTwoFactorTokenProvider<User>>();
         userTwoFactorTokenProvider
@@ -74,15 +77,23 @@ public class TwoFactorEmailServiceTests
             .GenerateAsync("TwoFactor", Arg.Any<UserManager<User>>(), user)
             .Returns(Task.FromResult(token));
 
-        var context = Substitute.For<ICurrentContext>();
-        context.DeviceType = DeviceType.Android;
-        context.IpAddress = IpAddress;
+        var userManager = sutProvider.GetDependency<UserManager<User>>();
+        userManager.RegisterTokenProvider(CoreHelpers.CustomProviderName(TwoFactorProviderType.Email), userTwoFactorTokenProvider);
+
+        user.SetTwoFactorProviders(new Dictionary<TwoFactorProviderType, TwoFactorProvider>
+        {
+            [TwoFactorProviderType.Email] = new TwoFactorProvider
+            {
+                MetaData = new Dictionary<string, object> { ["Email"] = email },
+                Enabled = true
+            }
+        });
 
         await sutProvider.Sut.SendTwoFactorEmailAsync(user);
 
         await sutProvider.GetDependency<IMailService>()
             .Received(1)
-            .SendTwoFactorEmailAsync(email, user.Email, token, IpAddress, deviceType,
+            .SendTwoFactorEmailAsync(email, user.Email, token, IpAddress, deviceType.ToString(),
                 TwoFactorEmailPurpose.Login);
     }
 
@@ -92,28 +103,28 @@ public class TwoFactorEmailServiceTests
         var email = user.Email.ToLowerInvariant();
         var token = "thisisatokentocompare";
         var IpAddress = "1.1.1.1";
-        var deviceType = "Android";
+        var deviceType = DeviceType.Android;
+
+        var context = sutProvider.GetDependency<ICurrentContext>();
+        context.DeviceType = deviceType;
+        context.IpAddress = IpAddress;
 
         var userTwoFactorTokenProvider = Substitute.For<IUserTwoFactorTokenProvider<User>>();
         userTwoFactorTokenProvider
             .CanGenerateTwoFactorTokenAsync(Arg.Any<UserManager<User>>(), user)
             .Returns(Task.FromResult(true));
         userTwoFactorTokenProvider
-            .GenerateAsync("TwoFactor", Arg.Any<UserManager<User>>(), user)
+            .GenerateAsync("otp:" + user.Email, Arg.Any<UserManager<User>>(), user)
             .Returns(Task.FromResult(token));
 
-        var context = Substitute.For<ICurrentContext>();
-        context.DeviceType = DeviceType.Android;
-        context.IpAddress = IpAddress;
-
         var userManager = sutProvider.GetDependency<UserManager<User>>();
-        userManager.RegisterTokenProvider(CoreHelpers.CustomProviderName(TwoFactorProviderType.Email), userTwoFactorTokenProvider);
+        userManager.RegisterTokenProvider(TokenOptions.DefaultEmailProvider, userTwoFactorTokenProvider);
 
         await sutProvider.Sut.SendNewDeviceVerificationEmailAsync(user);
 
         await sutProvider.GetDependency<IMailService>()
             .Received(1)
-            .SendTwoFactorEmailAsync(email, user.Email, token, IpAddress, deviceType,
+            .SendTwoFactorEmailAsync(email, user.Email, token, IpAddress, deviceType.ToString(),
                 TwoFactorEmailPurpose.NewDeviceVerification);
     }
 
@@ -164,13 +175,18 @@ public class TwoFactorEmailServiceTests
     [Theory]
     [BitAutoData(DeviceType.UnknownBrowser, "Unknown Browser")]
     [BitAutoData(DeviceType.Android, "Android")]
-    public async Task SendTwoFactorEmailAsync_DeviceMatches(SutProvider<TwoFactorEmailService> sutProvider, DeviceType deviceType, string deviceTypeName,
+    public async Task SendTwoFactorEmailAsync_DeviceMatches(DeviceType deviceType, string deviceTypeName,
+        SutProvider<TwoFactorEmailService> sutProvider,
         User user)
     {
         var email = user.Email.ToLowerInvariant();
         var token = "thisisatokentocompare";
         var IpAddress = "1.1.1.1";
 
+        var context = sutProvider.GetDependency<ICurrentContext>();
+        context.DeviceType = deviceType;
+        context.IpAddress = IpAddress;
+
         var userTwoFactorTokenProvider = Substitute.For<IUserTwoFactorTokenProvider<User>>();
         userTwoFactorTokenProvider
             .CanGenerateTwoFactorTokenAsync(Arg.Any<UserManager<User>>(), user)
@@ -179,10 +195,6 @@ public class TwoFactorEmailServiceTests
             .GenerateAsync("TwoFactor", Arg.Any<UserManager<User>>(), user)
             .Returns(Task.FromResult(token));
 
-        var context = Substitute.For<ICurrentContext>();
-        context.DeviceType = deviceType;
-        context.IpAddress = IpAddress;
-
         var userManager = sutProvider.GetDependency<UserManager<User>>();
         userManager.RegisterTokenProvider(CoreHelpers.CustomProviderName(TwoFactorProviderType.Email), userTwoFactorTokenProvider);
 
@@ -199,7 +211,7 @@ public class TwoFactorEmailServiceTests
 
         await sutProvider.GetDependency<IMailService>()
             .Received(1)
-            .SendTwoFactorEmailAsync(Arg.Any<string>(), Arg.Any<string>(), Arg.Any<string>(), Arg.Any<string>(), deviceTypeName, Arg.Any<TwoFactorEmailPurpose>());
+            .SendTwoFactorEmailAsync(Arg.Any<string>(), Arg.Any<string>(), Arg.Any<string>(), Arg.Any<string>(), deviceTypeName, TwoFactorEmailPurpose.Login);
     }
 
     [Theory, BitAutoData]