From 8ff820de17154490578a7eb2c1f2c7287b877ee4 Mon Sep 17 00:00:00 2001
From: Maciej Zieniuk <mzieniuk@bitwarden.com>
Date: Wed, 23 Oct 2024 14:29:09 +0100
Subject: [PATCH] PM-10600: Startup fix

---
 .../NotificationHubPushRegistrationService.cs          | 10 ++++++----
 .../NotificationHubPushRegistrationServiceTests.cs     |  4 +++-
 2 files changed, 9 insertions(+), 5 deletions(-)

diff --git a/src/Core/NotificationHub/NotificationHubPushRegistrationService.cs b/src/Core/NotificationHub/NotificationHubPushRegistrationService.cs
index 8a6142df4e..2ebf2692ea 100644
--- a/src/Core/NotificationHub/NotificationHubPushRegistrationService.cs
+++ b/src/Core/NotificationHub/NotificationHubPushRegistrationService.cs
@@ -4,6 +4,7 @@ using Bit.Core.Repositories;
 using Bit.Core.Services;
 using Bit.Core.Utilities;
 using Microsoft.Azure.NotificationHubs;
+using Microsoft.Extensions.DependencyInjection;
 
 namespace Bit.Core.NotificationHub;
 
@@ -11,16 +12,16 @@ public class NotificationHubPushRegistrationService : IPushRegistrationService
 {
     private readonly IInstallationDeviceRepository _installationDeviceRepository;
     private readonly INotificationHubPool _notificationHubPool;
-    private readonly IFeatureService _featureService;
+    private readonly IServiceProvider _serviceProvider;
 
     public NotificationHubPushRegistrationService(
         IInstallationDeviceRepository installationDeviceRepository,
         INotificationHubPool notificationHubPool,
-        IFeatureService featureService)
+        IServiceProvider serviceProvider)
     {
         _installationDeviceRepository = installationDeviceRepository;
         _notificationHubPool = notificationHubPool;
-        _featureService = featureService;
+        _serviceProvider = serviceProvider;
     }
 
     public async Task CreateOrUpdateRegistrationAsync(string pushToken, string deviceId, string userId,
@@ -51,7 +52,8 @@ public class NotificationHubPushRegistrationService : IPushRegistrationService
         switch (type)
         {
             case DeviceType.Android:
-                if (_featureService.IsEnabled(FeatureFlagKeys.AnhFcmv1Migration))
+                var featureService = _serviceProvider.GetRequiredService<IFeatureService>();
+                if (featureService.IsEnabled(FeatureFlagKeys.AnhFcmv1Migration))
                 {
                     payloadTemplate = "{\"message\":{\"data\":{\"type\":\"$(type)\",\"payload\":\"$(payload)\"}}}";
                     messageTemplate = "{\"message\":{\"data\":{\"type\":\"$(type)\"}," +
diff --git a/test/Core.Test/NotificationHub/NotificationHubPushRegistrationServiceTests.cs b/test/Core.Test/NotificationHub/NotificationHubPushRegistrationServiceTests.cs
index 00e94fe8f8..69b953333c 100644
--- a/test/Core.Test/NotificationHub/NotificationHubPushRegistrationServiceTests.cs
+++ b/test/Core.Test/NotificationHub/NotificationHubPushRegistrationServiceTests.cs
@@ -35,7 +35,9 @@ public class NotificationHubPushRegistrationServiceTests
     public async void CreateOrUpdateRegistrationAsync_DeviceTypeAndroid_InstallationCreated(bool identifierNull,
         SutProvider<NotificationHubPushRegistrationService> sutProvider, Guid deviceId, Guid userId, Guid? identifier)
     {
-        sutProvider.GetDependency<IFeatureService>().IsEnabled(FeatureFlagKeys.AnhFcmv1Migration).Returns(true);
+        var featureService = Substitute.For<IFeatureService>();
+        featureService.IsEnabled(FeatureFlagKeys.AnhFcmv1Migration).Returns(true);
+        sutProvider.GetDependency<IServiceProvider>().GetService(typeof(IFeatureService)).Returns(featureService);
         var notificationHubClient = Substitute.For<INotificationHubClient>();
         sutProvider.GetDependency<INotificationHubPool>().ClientFor(Arg.Any<Guid>()).Returns(notificationHubClient);