Refactor TwoFactorIsEnabledQuery to introduce VNextAsync methods for improved premium access checks and user detail handling. Removed obsolete feature service dependency and enhanced test coverage for new functionality.

2025-12-15 15:53:59 +00:00 · 2025-12-05 15:01:36 +00:00
parent 2d31237688
commit 123f57926f
3 changed files with 613 additions and 63 deletions
--- a/src/Core/Auth/UserFeatures/TwoFactorAuth/Interfaces/ITwoFactorIsEnabledQuery.cs
+++ b/src/Core/Auth/UserFeatures/TwoFactorAuth/Interfaces/ITwoFactorIsEnabledQuery.cs
@@ -22,4 +22,25 @@ public interface ITwoFactorIsEnabledQuery
    /// </summary>
    /// <param name="user">The user to check.</param>
    Task<bool> TwoFactorIsEnabledAsync(ITwoFactorProvidersUser user);
+
+    /// <summary>
+    /// Returns a list of user IDs and whether two factor is enabled for each user.
+    /// This version uses PremiumAccessQuery with cached organization abilities for better performance.
+    /// </summary>
+    /// <param name="userIds">The list of user IDs to check.</param>
+    Task<IEnumerable<(Guid userId, bool twoFactorIsEnabled)>> TwoFactorIsEnabledVNextAsync(IEnumerable<Guid> userIds);
+    /// <summary>
+    /// Returns a list of users and whether two factor is enabled for each user.
+    /// This version uses PremiumAccessQuery with cached organization abilities for better performance.
+    /// </summary>
+    /// <param name="users">The list of users to check.</param>
+    /// <typeparam name="T">The type of user in the list. Must implement <see cref="ITwoFactorProvidersUser"/>.</typeparam>
+    Task<IEnumerable<(T user, bool twoFactorIsEnabled)>> TwoFactorIsEnabledVNextAsync<T>(IEnumerable<T> users) where T : ITwoFactorProvidersUser;
+    /// <summary>
+    /// Returns whether two factor is enabled for the user. A user is able to have a TwoFactorProvider that is enabled but requires Premium.
+    /// If the user does not have premium then the TwoFactorProvider is considered _not_ enabled.
+    /// This version uses PremiumAccessQuery with cached organization abilities for better performance.
+    /// </summary>
+    /// <param name="user">The user to check.</param>
+    Task<bool> TwoFactorIsEnabledVNextAsync(ITwoFactorProvidersUser user);
 }
--- a/src/Core/Auth/UserFeatures/TwoFactorAuth/TwoFactorIsEnabledQuery.cs
+++ b/src/Core/Auth/UserFeatures/TwoFactorAuth/TwoFactorIsEnabledQuery.cs
@@ -8,7 +8,6 @@ using Bit.Core.Auth.UserFeatures.TwoFactorAuth.Interfaces;
 using Bit.Core.Entities;
 using Bit.Core.Models.Data.Organizations.OrganizationUsers;
 using Bit.Core.Repositories;
-using Bit.Core.Services;

 namespace Bit.Core.Auth.UserFeatures.TwoFactorAuth;

@@ -16,16 +15,13 @@ public class TwoFactorIsEnabledQuery : ITwoFactorIsEnabledQuery
 {
    private readonly IUserRepository _userRepository;
    private readonly IPremiumAccessQuery _premiumAccessQuery;
-    private readonly IFeatureService _featureService;

    public TwoFactorIsEnabledQuery(
        IUserRepository userRepository,
-        IPremiumAccessQuery premiumAccessQuery,
-        IFeatureService featureService)
+        IPremiumAccessQuery premiumAccessQuery)
    {
        _userRepository = userRepository;
        _premiumAccessQuery = premiumAccessQuery;
-        _featureService = featureService;
    }

    public async Task<IEnumerable<(Guid userId, bool twoFactorIsEnabled)>> TwoFactorIsEnabledAsync(IEnumerable<Guid> userIds)
@@ -36,34 +32,15 @@ public class TwoFactorIsEnabledQuery : ITwoFactorIsEnabledQuery
            return result;
        }

-        if (_featureService.IsEnabled(FeatureFlagKeys.PremiumAccessCacheCheck))
+        var userDetails = await _userRepository.GetManyWithCalculatedPremiumAsync([.. userIds]);
+        foreach (var userDetail in userDetails)
        {
-            var users = await _userRepository.GetManyAsync([.. userIds]);
-            var premiumStatus = await _premiumAccessQuery.CanAccessPremiumAsync(users);
-
-            foreach (var user in users)
-            {
-                result.Add(
-                    (user.Id,
-                     await TwoFactorEnabledAsync(
-                         user.GetTwoFactorProviders(),
-                         () => Task.FromResult(premiumStatus.GetValueOrDefault(user.Id, false))
-                     ))
-                );
-            }
-        }
-        else
-        {
-            var userDetails = await _userRepository.GetManyWithCalculatedPremiumAsync([.. userIds]);
-            foreach (var userDetail in userDetails)
-            {
-                result.Add(
-                    (userDetail.Id,
-                     await TwoFactorEnabledAsync(userDetail.GetTwoFactorProviders(),
-                                    () => Task.FromResult(userDetail.HasPremiumAccess))
-                    )
-                );
-            }
+            result.Add(
+                (userDetail.Id,
+                 await TwoFactorEnabledAsync(userDetail.GetTwoFactorProviders(),
+                                () => Task.FromResult(userDetail.HasPremiumAccess))
+                )
+            );
        }

        return result;
@@ -106,43 +83,102 @@ public class TwoFactorIsEnabledQuery : ITwoFactorIsEnabledQuery
            return false;
        }

-        if (_featureService.IsEnabled(FeatureFlagKeys.PremiumAccessCacheCheck))
+        return await TwoFactorEnabledAsync(
+            user.GetTwoFactorProviders(),
+            async () =>
+            {
+                var calcUser = await _userRepository.GetCalculatedPremiumAsync(userId.Value);
+                return calcUser?.HasPremiumAccess ?? false;
+            });
+    }
+
+    public async Task<IEnumerable<(Guid userId, bool twoFactorIsEnabled)>> TwoFactorIsEnabledVNextAsync(IEnumerable<Guid> userIds)
+    {
+        var result = new List<(Guid userId, bool hasTwoFactor)>();
+        if (userIds == null || !userIds.Any())
        {
-            // Try to get premium status without fetching User entity if possible
-            bool hasPersonalPremium;
-            if (user is User userEntity)
+            return result;
+        }
+
+        var users = await _userRepository.GetManyAsync([.. userIds]);
+        var premiumStatus = await _premiumAccessQuery.CanAccessPremiumAsync(users);
+
+        foreach (var user in users)
+        {
+            result.Add(
+                (user.Id,
+                 await TwoFactorEnabledAsync(
+                     user.GetTwoFactorProviders(),
+                     () => Task.FromResult(premiumStatus.GetValueOrDefault(user.Id, false))
+                 ))
+            );
+        }
+
+        return result;
+    }
+
+    public async Task<IEnumerable<(T user, bool twoFactorIsEnabled)>> TwoFactorIsEnabledVNextAsync<T>(IEnumerable<T> users)
+        where T : ITwoFactorProvidersUser
+    {
+        var userIds = users
+            .Select(u => u.GetUserId())
+            .Where(u => u.HasValue)
+            .Select(u => u.Value)
+            .ToList();
+
+        var twoFactorResults = await TwoFactorIsEnabledVNextAsync(userIds);
+
+        var result = new List<(T user, bool twoFactorIsEnabled)>();
+
+        foreach (var user in users)
+        {
+            var userId = user.GetUserId();
+            if (userId.HasValue)
            {
-                hasPersonalPremium = userEntity.Premium;
-            }
-            else if (user is OrganizationUserUserDetails orgUserDetails)
-            {
-                hasPersonalPremium = orgUserDetails.Premium.GetValueOrDefault(false);
+                var hasTwoFactor = twoFactorResults.FirstOrDefault(res => res.userId == userId.Value).twoFactorIsEnabled;
+                result.Add((user, hasTwoFactor));
            }
            else
            {
-                // Fallback: fetch the User entity
-                var fetchedUser = await _userRepository.GetByIdAsync(userId.Value);
-                if (fetchedUser == null)
-                {
-                    return false;
-                }
-                hasPersonalPremium = fetchedUser.Premium;
+                result.Add((user, false));
            }
+        }

-            return await TwoFactorEnabledAsync(
-                user.GetTwoFactorProviders(),
-                async () => await _premiumAccessQuery.CanAccessPremiumAsync(userId.Value, hasPersonalPremium));
+        return result;
+    }
+
+    public async Task<bool> TwoFactorIsEnabledVNextAsync(ITwoFactorProvidersUser user)
+    {
+        var userId = user.GetUserId();
+        if (!userId.HasValue)
+        {
+            return false;
+        }
+
+        // Try to get premium status without fetching User entity if possible
+        bool hasPersonalPremium;
+        if (user is User userEntity)
+        {
+            hasPersonalPremium = userEntity.Premium;
+        }
+        else if (user is OrganizationUserUserDetails orgUserDetails)
+        {
+            hasPersonalPremium = orgUserDetails.Premium.GetValueOrDefault(false);
        }
        else
        {
-            return await TwoFactorEnabledAsync(
-                user.GetTwoFactorProviders(),
-                async () =>
-                {
-                    var calcUser = await _userRepository.GetCalculatedPremiumAsync(userId.Value);
-                    return calcUser?.HasPremiumAccess ?? false;
-                });
+            // Fallback: fetch the User entity
+            var fetchedUser = await _userRepository.GetByIdAsync(userId.Value);
+            if (fetchedUser == null)
+            {
+                return false;
+            }
+            hasPersonalPremium = fetchedUser.Premium;
        }
+
+        return await TwoFactorEnabledAsync(
+            user.GetTwoFactorProviders(),
+            async () => await _premiumAccessQuery.CanAccessPremiumAsync(userId.Value, hasPersonalPremium));
    }

    /// <summary>
--- a/test/Core.Test/Auth/UserFeatures/TwoFactorAuth/TwoFactorIsEnabledQueryTests.cs
+++ b/test/Core.Test/Auth/UserFeatures/TwoFactorAuth/TwoFactorIsEnabledQueryTests.cs
@@ -1,5 +1,6 @@
 using Bit.Core.Auth.Enums;
 using Bit.Core.Auth.Models;
+using Bit.Core.Auth.UserFeatures.PremiumAccess;
 using Bit.Core.Auth.UserFeatures.TwoFactorAuth;
 using Bit.Core.Entities;
 using Bit.Core.Models.Data;
@@ -404,6 +405,503 @@ public class TwoFactorIsEnabledQueryTests
            .GetCalculatedPremiumAsync(default);
    }

+    [Theory]
+    [BitAutoData(TwoFactorProviderType.Authenticator)]
+    [BitAutoData(TwoFactorProviderType.Email)]
+    [BitAutoData(TwoFactorProviderType.Remember)]
+    [BitAutoData(TwoFactorProviderType.OrganizationDuo)]
+    [BitAutoData(TwoFactorProviderType.WebAuthn)]
+    public async Task TwoFactorIsEnabledVNextAsync_WithProviderTypeNotRequiringPremium_ReturnsAllTwoFactorEnabled(
+        TwoFactorProviderType freeProviderType,
+        SutProvider<TwoFactorIsEnabledQuery> sutProvider,
+        List<User> users)
+    {
+        // Arrange
+        var userIds = users.Select(u => u.Id).ToList();
+        var twoFactorProviders = new Dictionary<TwoFactorProviderType, TwoFactorProvider>
+        {
+            { freeProviderType, new TwoFactorProvider { Enabled = true } }
+        };
+
+        foreach (var user in users)
+        {
+            user.Premium = false;
+            user.SetTwoFactorProviders(twoFactorProviders);
+        }
+
+        var premiumStatus = users.ToDictionary(u => u.Id, u => false);
+
+        sutProvider.GetDependency<IUserRepository>()
+            .GetManyAsync(Arg.Is<IEnumerable<Guid>>(i => i.All(userIds.Contains)))
+            .Returns(users);
+
+        sutProvider.GetDependency<IPremiumAccessQuery>()
+            .CanAccessPremiumAsync(Arg.Is<IEnumerable<User>>(u => u.All(users.Contains)))
+            .Returns(premiumStatus);
+
+        // Act
+        var result = await sutProvider.Sut.TwoFactorIsEnabledVNextAsync(userIds);
+
+        // Assert
+        foreach (var user in users)
+        {
+            Assert.Contains(result, res => res.userId == user.Id && res.twoFactorIsEnabled == true);
+        }
+    }
+
+    [Theory, BitAutoData]
+    public async Task TwoFactorIsEnabledVNextAsync_DatabaseReturnsEmpty_ResultEmpty(
+        SutProvider<TwoFactorIsEnabledQuery> sutProvider,
+        List<User> users)
+    {
+        // Arrange
+        var userIds = users.Select(u => u.Id).ToList();
+
+        sutProvider.GetDependency<IUserRepository>()
+            .GetManyAsync(Arg.Any<IEnumerable<Guid>>())
+            .Returns([]);
+
+        sutProvider.GetDependency<IPremiumAccessQuery>()
+            .CanAccessPremiumAsync(Arg.Any<IEnumerable<User>>())
+            .Returns(new Dictionary<Guid, bool>());
+
+        // Act
+        var result = await sutProvider.Sut.TwoFactorIsEnabledVNextAsync(userIds);
+
+        // Assert
+        Assert.Empty(result);
+    }
+
+    [Theory]
+    [BitAutoData((IEnumerable<Guid>)null)]
+    [BitAutoData([])]
+    public async Task TwoFactorIsEnabledVNextAsync_UserIdsNullorEmpty_ResultEmpty(
+        IEnumerable<Guid> userIds,
+        SutProvider<TwoFactorIsEnabledQuery> sutProvider)
+    {
+        // Act
+        var result = await sutProvider.Sut.TwoFactorIsEnabledVNextAsync(userIds);
+
+        // Assert
+        Assert.Empty(result);
+    }
+
+    [Theory]
+    [BitAutoData]
+    public async Task TwoFactorIsEnabledVNextAsync_WithNoTwoFactorEnabled_ReturnsAllTwoFactorDisabled(
+        SutProvider<TwoFactorIsEnabledQuery> sutProvider,
+        List<User> users)
+    {
+        // Arrange
+        var userIds = users.Select(u => u.Id).ToList();
+        var twoFactorProviders = new Dictionary<TwoFactorProviderType, TwoFactorProvider>
+        {
+            { TwoFactorProviderType.Email, new TwoFactorProvider { Enabled = false } }
+        };
+
+        foreach (var user in users)
+        {
+            user.SetTwoFactorProviders(twoFactorProviders);
+        }
+
+        var premiumStatus = users.ToDictionary(u => u.Id, u => false);
+
+        sutProvider.GetDependency<IUserRepository>()
+            .GetManyAsync(Arg.Is<IEnumerable<Guid>>(i => i.All(userIds.Contains)))
+            .Returns(users);
+
+        sutProvider.GetDependency<IPremiumAccessQuery>()
+            .CanAccessPremiumAsync(Arg.Is<IEnumerable<User>>(u => u.All(users.Contains)))
+            .Returns(premiumStatus);
+
+        // Act
+        var result = await sutProvider.Sut.TwoFactorIsEnabledVNextAsync(userIds);
+
+        // Assert
+        foreach (var user in users)
+        {
+            Assert.Contains(result, res => res.userId == user.Id && res.twoFactorIsEnabled == false);
+        }
+    }
+
+    [Theory]
+    [BitAutoData(TwoFactorProviderType.Duo)]
+    [BitAutoData(TwoFactorProviderType.YubiKey)]
+    public async Task TwoFactorIsEnabledVNextAsync_WithProviderTypeRequiringPremium_ReturnsMixedResults(
+        TwoFactorProviderType premiumProviderType,
+        SutProvider<TwoFactorIsEnabledQuery> sutProvider,
+        List<User> users)
+    {
+        // Arrange
+        var userIds = users.Select(u => u.Id).ToList();
+        var twoFactorProviders = new Dictionary<TwoFactorProviderType, TwoFactorProvider>
+        {
+            { TwoFactorProviderType.Email, new TwoFactorProvider { Enabled = false } },
+            { premiumProviderType, new TwoFactorProvider { Enabled = true } }
+        };
+
+        foreach (var user in users)
+        {
+            user.Premium = false;
+            user.SetTwoFactorProviders(twoFactorProviders);
+        }
+
+        // Only the first user has premium access
+        var premiumStatus = users.ToDictionary(
+            u => u.Id,
+            u => users.IndexOf(u) == 0);
+
+        sutProvider.GetDependency<IUserRepository>()
+            .GetManyAsync(Arg.Is<IEnumerable<Guid>>(i => i.All(userIds.Contains)))
+            .Returns(users);
+
+        sutProvider.GetDependency<IPremiumAccessQuery>()
+            .CanAccessPremiumAsync(Arg.Is<IEnumerable<User>>(u => u.All(users.Contains)))
+            .Returns(premiumStatus);
+
+        // Act
+        var result = await sutProvider.Sut.TwoFactorIsEnabledVNextAsync(userIds);
+
+        // Assert
+        foreach (var user in users)
+        {
+            var expectedEnabled = premiumStatus[user.Id];
+            Assert.Contains(result, res => res.userId == user.Id && res.twoFactorIsEnabled == expectedEnabled);
+        }
+    }
+
+    [Theory]
+    [BitAutoData("")]
+    [BitAutoData("{}")]
+    [BitAutoData((string)null)]
+    public async Task TwoFactorIsEnabledVNextAsync_WithNullOrEmptyTwoFactorProviders_ReturnsAllTwoFactorDisabled(
+        string twoFactorProviders,
+        SutProvider<TwoFactorIsEnabledQuery> sutProvider,
+        List<User> users)
+    {
+        // Arrange
+        var userIds = users.Select(u => u.Id).ToList();
+
+        foreach (var user in users)
+        {
+            user.TwoFactorProviders = twoFactorProviders;
+        }
+
+        var premiumStatus = users.ToDictionary(u => u.Id, u => false);
+
+        sutProvider.GetDependency<IUserRepository>()
+            .GetManyAsync(Arg.Is<IEnumerable<Guid>>(i => i.All(userIds.Contains)))
+            .Returns(users);
+
+        sutProvider.GetDependency<IPremiumAccessQuery>()
+            .CanAccessPremiumAsync(Arg.Is<IEnumerable<User>>(u => u.All(users.Contains)))
+            .Returns(premiumStatus);
+
+        // Act
+        var result = await sutProvider.Sut.TwoFactorIsEnabledVNextAsync(userIds);
+
+        // Assert
+        foreach (var user in users)
+        {
+            Assert.Contains(result, res => res.userId == user.Id && res.twoFactorIsEnabled == false);
+        }
+    }
+
+    [Theory]
+    [BitAutoData]
+    public async Task TwoFactorIsEnabledVNextAsync_Generic_WithNoUserIds_ReturnsAllTwoFactorDisabled(
+        SutProvider<TwoFactorIsEnabledQuery> sutProvider,
+        List<OrganizationUserUserDetails> users)
+    {
+        // Arrange
+        foreach (var user in users)
+        {
+            user.UserId = null;
+        }
+
+        // Act
+        var result = await sutProvider.Sut.TwoFactorIsEnabledVNextAsync(users);
+
+        // Assert
+        foreach (var user in users)
+        {
+            Assert.Contains(result, res => res.user.Equals(user) && res.twoFactorIsEnabled == false);
+        }
+
+        // No UserIds were supplied so no calls to the UserRepository should have been made
+        await sutProvider.GetDependency<IUserRepository>()
+            .DidNotReceiveWithAnyArgs()
+            .GetManyAsync(default);
+    }
+
+    [Theory]
+    [BitAutoData]
+    public async Task TwoFactorIsEnabledVNextAsync_SingleUser_UserIdNull_ReturnsFalse(
+        SutProvider<TwoFactorIsEnabledQuery> sutProvider)
+    {
+        // Arrange
+        var user = new TestTwoFactorProviderUser
+        {
+            Id = null
+        };
+
+        // Act
+        var result = await sutProvider.Sut.TwoFactorIsEnabledVNextAsync(user);
+
+        // Assert
+        Assert.False(result);
+    }
+
+    [Theory]
+    [BitAutoData(TwoFactorProviderType.Authenticator)]
+    [BitAutoData(TwoFactorProviderType.Email)]
+    [BitAutoData(TwoFactorProviderType.Remember)]
+    [BitAutoData(TwoFactorProviderType.OrganizationDuo)]
+    [BitAutoData(TwoFactorProviderType.WebAuthn)]
+    public async Task TwoFactorIsEnabledVNextAsync_SingleUser_WithProviderTypeNotRequiringPremium_ReturnsTrue(
+        TwoFactorProviderType freeProviderType,
+        SutProvider<TwoFactorIsEnabledQuery> sutProvider,
+        User user)
+    {
+        // Arrange
+        var twoFactorProviders = new Dictionary<TwoFactorProviderType, TwoFactorProvider>
+        {
+            { freeProviderType, new TwoFactorProvider { Enabled = true } }
+        };
+
+        user.Premium = false;
+        user.SetTwoFactorProviders(twoFactorProviders);
+
+        // Act
+        var result = await sutProvider.Sut.TwoFactorIsEnabledVNextAsync(user);
+
+        // Assert
+        Assert.True(result);
+
+        // Should not need to check premium access for free providers
+        await sutProvider.GetDependency<IPremiumAccessQuery>()
+            .DidNotReceiveWithAnyArgs()
+            .CanAccessPremiumAsync(default(Guid), default);
+    }
+
+    [Theory]
+    [BitAutoData]
+    public async Task TwoFactorIsEnabledVNextAsync_SingleUser_WithNoTwoFactorEnabled_ReturnsFalse(
+        SutProvider<TwoFactorIsEnabledQuery> sutProvider,
+        User user)
+    {
+        // Arrange
+        var twoFactorProviders = new Dictionary<TwoFactorProviderType, TwoFactorProvider>
+        {
+            { TwoFactorProviderType.Email, new TwoFactorProvider { Enabled = false } }
+        };
+
+        user.SetTwoFactorProviders(twoFactorProviders);
+
+        // Act
+        var result = await sutProvider.Sut.TwoFactorIsEnabledVNextAsync(user);
+
+        // Assert
+        Assert.False(result);
+
+        await sutProvider.GetDependency<IPremiumAccessQuery>()
+            .DidNotReceiveWithAnyArgs()
+            .CanAccessPremiumAsync(default(Guid), default);
+    }
+
+    [Theory]
+    [BitAutoData(TwoFactorProviderType.Duo)]
+    [BitAutoData(TwoFactorProviderType.YubiKey)]
+    public async Task TwoFactorIsEnabledVNextAsync_SingleUser_WithProviderTypeRequiringPremium_WithoutPremium_ReturnsFalse(
+        TwoFactorProviderType premiumProviderType,
+        SutProvider<TwoFactorIsEnabledQuery> sutProvider,
+        User user)
+    {
+        // Arrange
+        var twoFactorProviders = new Dictionary<TwoFactorProviderType, TwoFactorProvider>
+        {
+            { premiumProviderType, new TwoFactorProvider { Enabled = true } }
+        };
+
+        user.Premium = false;
+        user.SetTwoFactorProviders(twoFactorProviders);
+
+        sutProvider.GetDependency<IPremiumAccessQuery>()
+            .CanAccessPremiumAsync(user.Id, false)
+            .Returns(false);
+
+        // Act
+        var result = await sutProvider.Sut.TwoFactorIsEnabledVNextAsync(user);
+
+        // Assert
+        Assert.False(result);
+
+        await sutProvider.GetDependency<IPremiumAccessQuery>()
+            .Received(1)
+            .CanAccessPremiumAsync(user.Id, false);
+    }
+
+    [Theory]
+    [BitAutoData(TwoFactorProviderType.Duo)]
+    [BitAutoData(TwoFactorProviderType.YubiKey)]
+    public async Task TwoFactorIsEnabledVNextAsync_SingleUser_WithProviderTypeRequiringPremium_WithPersonalPremium_ReturnsTrue(
+        TwoFactorProviderType premiumProviderType,
+        SutProvider<TwoFactorIsEnabledQuery> sutProvider,
+        User user)
+    {
+        // Arrange
+        var twoFactorProviders = new Dictionary<TwoFactorProviderType, TwoFactorProvider>
+        {
+            { premiumProviderType, new TwoFactorProvider { Enabled = true } }
+        };
+
+        user.Premium = true;
+        user.SetTwoFactorProviders(twoFactorProviders);
+
+        sutProvider.GetDependency<IPremiumAccessQuery>()
+            .CanAccessPremiumAsync(user.Id, true)
+            .Returns(true);
+
+        // Act
+        var result = await sutProvider.Sut.TwoFactorIsEnabledVNextAsync(user);
+
+        // Assert
+        Assert.True(result);
+
+        await sutProvider.GetDependency<IPremiumAccessQuery>()
+            .Received(1)
+            .CanAccessPremiumAsync(user.Id, true);
+    }
+
+    [Theory]
+    [BitAutoData(TwoFactorProviderType.Duo)]
+    [BitAutoData(TwoFactorProviderType.YubiKey)]
+    public async Task TwoFactorIsEnabledVNextAsync_SingleUser_WithProviderTypeRequiringPremium_WithOrgPremium_ReturnsTrue(
+        TwoFactorProviderType premiumProviderType,
+        SutProvider<TwoFactorIsEnabledQuery> sutProvider,
+        User user)
+    {
+        // Arrange
+        var twoFactorProviders = new Dictionary<TwoFactorProviderType, TwoFactorProvider>
+        {
+            { premiumProviderType, new TwoFactorProvider { Enabled = true } }
+        };
+
+        user.Premium = false;
+        user.SetTwoFactorProviders(twoFactorProviders);
+
+        sutProvider.GetDependency<IPremiumAccessQuery>()
+            .CanAccessPremiumAsync(user.Id, false)
+            .Returns(true); // Has premium from org
+
+        // Act
+        var result = await sutProvider.Sut.TwoFactorIsEnabledVNextAsync(user);
+
+        // Assert
+        Assert.True(result);
+
+        await sutProvider.GetDependency<IPremiumAccessQuery>()
+            .Received(1)
+            .CanAccessPremiumAsync(user.Id, false);
+    }
+
+    [Theory]
+    [BitAutoData]
+    public async Task TwoFactorIsEnabledVNextAsync_SingleUser_WithNullTwoFactorProviders_ReturnsFalse(
+        SutProvider<TwoFactorIsEnabledQuery> sutProvider,
+        User user)
+    {
+        // Arrange
+        user.TwoFactorProviders = null;
+
+        // Act
+        var result = await sutProvider.Sut.TwoFactorIsEnabledVNextAsync(user);
+
+        // Assert
+        Assert.False(result);
+        await sutProvider.GetDependency<IPremiumAccessQuery>()
+            .DidNotReceiveWithAnyArgs()
+            .CanAccessPremiumAsync(default(Guid), default);
+    }
+
+    [Theory]
+    [BitAutoData(TwoFactorProviderType.Duo)]
+    [BitAutoData(TwoFactorProviderType.YubiKey)]
+    public async Task TwoFactorIsEnabledVNextAsync_SingleUser_OrganizationUserUserDetails_WithPremium_ReturnsTrue(
+        TwoFactorProviderType premiumProviderType,
+        SutProvider<TwoFactorIsEnabledQuery> sutProvider,
+        OrganizationUserUserDetails orgUserDetails)
+    {
+        // Arrange
+        var twoFactorProviders = new Dictionary<TwoFactorProviderType, TwoFactorProvider>
+        {
+            { premiumProviderType, new TwoFactorProvider { Enabled = true } }
+        };
+
+        orgUserDetails.Premium = false;
+        orgUserDetails.TwoFactorProviders = JsonHelpers.LegacySerialize(twoFactorProviders, JsonHelpers.LegacyEnumKeyResolver);
+
+        sutProvider.GetDependency<IPremiumAccessQuery>()
+            .CanAccessPremiumAsync(orgUserDetails.UserId!.Value, false)
+            .Returns(true);
+
+        // Act
+        var result = await sutProvider.Sut.TwoFactorIsEnabledVNextAsync(orgUserDetails);
+
+        // Assert
+        Assert.True(result);
+
+        await sutProvider.GetDependency<IPremiumAccessQuery>()
+            .Received(1)
+            .CanAccessPremiumAsync(orgUserDetails.UserId.Value, false);
+    }
+
+    [Theory]
+    [BitAutoData(TwoFactorProviderType.Duo)]
+    [BitAutoData(TwoFactorProviderType.YubiKey)]
+    public async Task TwoFactorIsEnabledVNextAsync_SingleUser_UnknownType_FetchesUser(
+        TwoFactorProviderType premiumProviderType,
+        SutProvider<TwoFactorIsEnabledQuery> sutProvider,
+        User fetchedUser)
+    {
+        // Arrange
+        var twoFactorProviders = new Dictionary<TwoFactorProviderType, TwoFactorProvider>
+        {
+            { premiumProviderType, new TwoFactorProvider { Enabled = true } }
+        };
+
+        var testUser = new TestTwoFactorProviderUser
+        {
+            Id = fetchedUser.Id,
+            Premium = false,
+            TwoFactorProviders = JsonHelpers.LegacySerialize(twoFactorProviders, JsonHelpers.LegacyEnumKeyResolver)
+        };
+
+        fetchedUser.Premium = false;
+
+        sutProvider.GetDependency<IUserRepository>()
+            .GetByIdAsync(fetchedUser.Id)
+            .Returns(fetchedUser);
+
+        sutProvider.GetDependency<IPremiumAccessQuery>()
+            .CanAccessPremiumAsync(fetchedUser.Id, false)
+            .Returns(true);
+
+        // Act
+        var result = await sutProvider.Sut.TwoFactorIsEnabledVNextAsync(testUser);
+
+        // Assert
+        Assert.True(result);
+
+        await sutProvider.GetDependency<IUserRepository>()
+            .Received(1)
+            .GetByIdAsync(fetchedUser.Id);
+
+        await sutProvider.GetDependency<IPremiumAccessQuery>()
+            .Received(1)
+            .CanAccessPremiumAsync(fetchedUser.Id, false);
+    }
+
    private class TestTwoFactorProviderUser : ITwoFactorProvidersUser
    {
        public Guid? Id { get; set; }
@@ -418,10 +916,5 @@ public class TwoFactorIsEnabledQueryTests
        {
            return Id;
        }
-
-        public bool GetPremium()
-        {
-            return Premium;
-        }
    }
 }