DEV2-3003: Fix not listening for capabilities changes (#561)

* Fix bug for not updating of capabilities change * Formatting
codota · Jul 9, 2023 · 492ba7f · 492ba7f
1 parent ba87d92
commit 492ba7f
Show file tree

Hide file tree

Showing 8 changed files with 69 additions and 89 deletions.
diff --git a/Common/src/main/java/com/tabnineCommon/capabilities/Capabilities.kt b/Common/src/main/java/com/tabnineCommon/capabilities/Capabilities.kt
@@ -0,0 +1,12 @@
+package com.tabnineCommon.capabilities
+
+import com.tabnineCommon.binary.requests.capabilities.ExperimentSource
+
+data class Capabilities(
+    val features: Set<Capability>,
+    val experimentSource: ExperimentSource?
+) {
+    fun isReady() = experimentSource == null || experimentSource.isRemoteBasedSource()
+
+    fun isEnabled(capability: Capability) = features.contains(capability)
+}
diff --git a/Common/src/main/java/com/tabnineCommon/capabilities/CapabilitiesService.java b/Common/src/main/java/com/tabnineCommon/capabilities/CapabilitiesService.java
@@ -14,7 +14,6 @@
 import java.util.Objects;
 import java.util.Optional;
 import java.util.Set;
-import java.util.concurrent.atomic.AtomicBoolean;
 
 public class CapabilitiesService {
 
@@ -28,7 +27,6 @@ public class CapabilitiesService {
   private final BinaryRequestFacade binaryRequestFacade =
       DependencyContainer.instanceOfBinaryRequestFacade();
   private final Set<Capability> enabledCapabilities = new HashSet<>();
-  private final AtomicBoolean isRemoteBasedSource = new AtomicBoolean(false);
 
   public static CapabilitiesService getInstance() {
     return ServiceManager.getService(CapabilitiesService.class);
@@ -38,10 +36,6 @@ public void init() {
     scheduleFetchCapabilitiesTask();
   }
 
-  public boolean isReady() {
-    return this.isRemoteBasedSource.get();
-  }
-
   public boolean isCapabilityEnabled(Capability capability) {
     if (Config.IS_SELF_HOSTED) {
       return true;
@@ -107,11 +101,6 @@ private void fetchCapabilities() {
     if (capabilitiesResponse.getEnabledFeatures() != null) {
       setCapabilities(capabilitiesResponse);
     }
-
-    if (capabilitiesResponse.getExperimentSource() == null
-        || capabilitiesResponse.getExperimentSource().isRemoteBasedSource()) {
-      isRemoteBasedSource.set(true);
-    }
   }
 
   private void setCapabilities(CapabilitiesResponse capabilitiesResponse) {
@@ -125,7 +114,8 @@ private void setCapabilities(CapabilitiesResponse capabilitiesResponse) {
       if (!newCapabilities.equals(enabledCapabilities)) {
         enabledCapabilities.clear();
         enabledCapabilities.addAll(newCapabilities);
-        CapabilityNotifier.Companion.publish(enabledCapabilities);
+        CapabilityNotifier.Companion.publish(
+            new Capabilities(newCapabilities, capabilitiesResponse.getExperimentSource()));
       }
     }
   }

diff --git a/Common/src/main/java/com/tabnineCommon/capabilities/CapabilityNotifier.kt b/Common/src/main/java/com/tabnineCommon/capabilities/CapabilityNotifier.kt
@@ -7,7 +7,7 @@ import com.intellij.util.messages.Topic
 fun interface CapabilityNotifier {
     companion object {
         private val CAPABILITY_CHANGED_TOPIC = Topic.create("com.tabnine.capabilties", CapabilityNotifier::class.java)
-        fun publish(state: Set<Capability>) {
+        fun publish(state: Capabilities) {
             ApplicationManager.getApplication()?.messageBus
                 ?.syncPublisher(CAPABILITY_CHANGED_TOPIC)
                 ?.stateChanged(state)
@@ -18,5 +18,5 @@ fun interface CapabilityNotifier {
             return bus
         }
     }
-    fun stateChanged(state: Set<Capability>)
+    fun stateChanged(state: Capabilities)
 }
diff --git a/Tabnine/src/main/java/com/tabnine/lifecycle/pushtosignin/PushToSignIn.kt b/Tabnine/src/main/java/com/tabnine/lifecycle/pushtosignin/PushToSignIn.kt
@@ -18,7 +18,7 @@ class PushToSignIn {
         }
         CapabilityNotifier.subscribe(
             CapabilityNotifier { state ->
-                if (state.contains(Capability.FORCE_REGISTRATION)) {
+                if (state.isEnabled(Capability.FORCE_REGISTRATION)) {
                     transition()
                 }
             }

diff --git a/Tabnine/src/main/java/com/tabnine/statusBar/TabnineIconProvider.kt b/Tabnine/src/main/java/com/tabnine/statusBar/TabnineIconProvider.kt
@@ -1,8 +1,6 @@
 package com.tabnine.statusBar
 
 import com.tabnineCommon.binary.requests.config.CloudConnectionHealthStatus
-import com.tabnineCommon.capabilities.CapabilitiesService
-import com.tabnineCommon.capabilities.Capability
 import com.tabnineCommon.general.ServiceLevel
 import com.tabnineCommon.general.StaticConfig
 import com.tabnineCommon.general.getSubscriptionType
@@ -13,12 +11,12 @@ object TabnineIconProvider {
     fun getIcon(
         serviceLevel: ServiceLevel?,
         isLoggedIn: Boolean?,
-        cloudConnectionHealthStatus: CloudConnectionHealthStatus
+        cloudConnectionHealthStatus: CloudConnectionHealthStatus,
+        isForcedRegistration: Boolean?
     ): Icon {
-        if (!CapabilitiesService.getInstance().isReady || isLoggedIn == null || serviceLevel == null || (
-            CapabilitiesService.getInstance()
-                .isCapabilityEnabled(Capability.FORCE_REGISTRATION) &&
-                !isLoggedIn && serviceLevel == ServiceLevel.FREE
+        if (isForcedRegistration == null || isLoggedIn == null || serviceLevel == null || (
+            isForcedRegistration &&
+                !isLoggedIn
             )
         ) {
             return StaticConfig.ICON_AND_NAME

diff --git a/Tabnine/src/main/java/com/tabnine/statusBar/TabnineStatusBarWidget.java b/Tabnine/src/main/java/com/tabnine/statusBar/TabnineStatusBarWidget.java
@@ -15,6 +15,7 @@
 import com.tabnineCommon.binary.requests.config.CloudConnectionHealthStatus;
 import com.tabnineCommon.capabilities.CapabilitiesService;
 import com.tabnineCommon.capabilities.Capability;
+import com.tabnineCommon.capabilities.CapabilityNotifier;
 import com.tabnineCommon.general.ServiceLevel;
 import com.tabnineCommon.lifecycle.BinaryStateChangeNotifier;
 import java.awt.event.MouseEvent;
@@ -32,6 +33,8 @@ public class TabnineStatusBarWidget extends EditorBasedWidget
   private ServiceLevel serviceLevel = null;
   private CloudConnectionHealthStatus cloudConnectionHealthStatus = CloudConnectionHealthStatus.Ok;
 
+  @Nullable private Boolean isForcedRegistration = null;
+
   public TabnineStatusBarWidget(@NotNull Project project) {
     super(project);
     // register for state changes (we will get notified whenever the state changes)
@@ -55,11 +58,26 @@ public TabnineStatusBarWidget(@NotNull Project project) {
               this.isLimited = limited;
               update();
             });
+
+    CapabilityNotifier.Companion.subscribe(
+        capabilities -> {
+          if (capabilities.isReady()) {
+            boolean newForceRegistration = capabilities.isEnabled(Capability.FORCE_REGISTRATION);
+
+            if (isForcedRegistration == null || isForcedRegistration != newForceRegistration) {
+              isForcedRegistration = newForceRegistration;
+              update();
+            }
+          }
+        });
   }
 
   public Icon getIcon() {
     return TabnineIconProvider.getIcon(
-        this.serviceLevel, this.isLoggedIn, this.cloudConnectionHealthStatus);
+        this.serviceLevel,
+        this.isLoggedIn,
+        this.cloudConnectionHealthStatus,
+        this.isForcedRegistration);
   }
 
   public @Nullable("null means the widget is unable to show the popup") ListPopup getPopupStep() {

diff --git a/Tabnine/src/test/java/com/tabnine/healthCheck/StatusBarWidgetTests.java b/Tabnine/src/test/java/com/tabnine/healthCheck/StatusBarWidgetTests.java
@@ -1,47 +1,56 @@
 package com.tabnine.healthCheck;
 
+import com.google.common.collect.Sets;
 import com.tabnine.MockedBinaryCompletionTestCase;
 import com.tabnine.statusBar.TabnineStatusBarWidget;
 import com.tabnine.testUtils.HealthCheckTestUtils;
-import com.tabnine.testUtils.MockBinaryResponse;
+import com.tabnineCommon.binary.requests.capabilities.ExperimentSource;
 import com.tabnineCommon.binary.requests.config.CloudConnectionHealthStatus;
-import com.tabnineCommon.capabilities.CapabilitiesService;
+import com.tabnineCommon.capabilities.Capabilities;
+import com.tabnineCommon.capabilities.Capability;
+import com.tabnineCommon.capabilities.CapabilityNotifier;
 import com.tabnineCommon.general.ServiceLevel;
 import com.tabnineCommon.general.StaticConfig;
-import java.util.ArrayList;
+import java.util.HashSet;
 import org.junit.Test;
 
 public class StatusBarWidgetTests extends MockedBinaryCompletionTestCase {
 
   @Test
-  public void should_get_connection_healthy_icon_when_connection_healthy()
-      throws InterruptedException {
+  public void should_get_connection_healthy_icon_when_connection_healthy() {
     TabnineStatusBarWidget widget = new TabnineStatusBarWidget(myFixture.getProject());
     HealthCheckTestUtils.notifyStateForWidget(
         ServiceLevel.FREE, true, CloudConnectionHealthStatus.Ok);
-    MockBinaryResponse.mockCapabilities(binaryProcessGatewayMock, "API", new ArrayList<>());
-    CapabilitiesService.getInstance().init();
-
-    while (!CapabilitiesService.getInstance().isReady()) {
-      Thread.sleep(100);
-    }
+    CapabilityNotifier.Companion.publish(new Capabilities(new HashSet<>(), ExperimentSource.API));
 
     assertEquals(StaticConfig.ICON_AND_NAME_STARTER, widget.getIcon());
   }
 
   @Test
-  public void should_get_connection_unhealthy_icon_when_connection_unhealthy()
-      throws InterruptedException {
+  public void should_get_connection_unhealthy_icon_when_connection_unhealthy() {
     TabnineStatusBarWidget widget = new TabnineStatusBarWidget(myFixture.getProject());
     HealthCheckTestUtils.notifyStateForWidget(
         ServiceLevel.FREE, true, CloudConnectionHealthStatus.Failed);
-    MockBinaryResponse.mockCapabilities(binaryProcessGatewayMock, "API", new ArrayList<>());
-    CapabilitiesService.getInstance().init();
-
-    while (!CapabilitiesService.getInstance().isReady()) {
-      Thread.sleep(100);
-    }
+    CapabilityNotifier.Companion.publish(new Capabilities(new HashSet<>(), ExperimentSource.API));
 
     assertEquals(StaticConfig.ICON_AND_NAME_CONNECTION_LOST_STARTER, widget.getIcon());
   }
+
+  @Test
+  public void
+      given_not_ready_source_then_ready_source_when_get_icon_should_return_no_suffix_then_suffix() {
+    TabnineStatusBarWidget widget = new TabnineStatusBarWidget(myFixture.getProject());
+    HealthCheckTestUtils.notifyStateForWidget(
+        ServiceLevel.FREE, true, CloudConnectionHealthStatus.Ok);
+    CapabilityNotifier.Companion.publish(
+        new Capabilities(new HashSet<>(), ExperimentSource.Unknown));
+
+    assertEquals(StaticConfig.ICON_AND_NAME, widget.getIcon());
+
+    CapabilityNotifier.Companion.publish(
+        new Capabilities(
+            Sets.immutableEnumSet(Capability.FORCE_REGISTRATION), ExperimentSource.API));
+
+    assertEquals(StaticConfig.ICON_AND_NAME_STARTER, widget.getIcon());
+  }
 }
diff --git a/Tabnine/src/test/java/com/tabnine/testUtils/MockBinaryResponse.kt b/Tabnine/src/test/java/com/tabnine/testUtils/MockBinaryResponse.kt