Internal change

src/java/com/google/devtools/mobileharness/fe/v6/service/config/util/ConfigServiceCapability.java

@@ -23,6 +23,7 @@
 import com.google.devtools.mobileharness.fe.v6.service.proto.config.HostPropertiesUiStatus;
 import com.google.devtools.mobileharness.fe.v6.service.proto.config.PartStatus;
 import com.google.devtools.mobileharness.fe.v6.service.util.Environment;
+import com.google.devtools.mobileharness.fe.v6.service.util.UniverseScope;
 import com.google.devtools.mobileharness.shared.util.flags.Flags;
 import com.google.inject.assistedinject.Assisted;
 import com.google.inject.assistedinject.AssistedInject;
@@ -34,26 +35,38 @@
 public class ConfigServiceCapability {
 
   private final Environment environment;
-  private final String universe;
+  private final UniverseScope universe;
 
   @AssistedInject
-  ConfigServiceCapability(Environment environment, @Assisted String universe) {
+  ConfigServiceCapability(Environment environment, @Assisted UniverseScope universe) {
     this.environment = environment;
     this.universe = universe;
   }
 
+  /**
+   * @deprecated Use {@link #ConfigServiceCapability(Environment, UniverseScope)} instead. TODO:
+   *     Remove after all callers are migrated to UniverseScope.
+   */
+  @Deprecated
+  @AssistedInject
+  ConfigServiceCapability(Environment environment, @Assisted String universe) {
+    this(environment, UniverseScope.fromString(universe));
+  }
+
+  // TODO: Consolidate with FeatureManager.isConfigurationFeatureEnabled() which
+  // contains similar universe-based availability logic.
   /** Returns true if the configuration service is available for the current context. */
   public boolean isConfigServiceAvailable() {
     if (environment.isGoogleInternal()) {
-      return universe.isEmpty() || universe.equals("google_1p");
+      return universe instanceof UniverseScope.SelfUniverse;
     }
 
     if (!Flags.instance().feConnectToConfigServer.get()) {
       return false;
     }
 
     // OSS/ATS only supports the default context (empty universe).
-    return universe.isEmpty();
+    return universe instanceof UniverseScope.SelfUniverse;
   }
 
   /**
@@ -68,7 +81,10 @@ public void checkConfigServiceAvailability() {
     if (environment.isGoogleInternal()) {
       throw new UnsupportedOperationException(
           String.format(
-              "Configuration operations are not currently supported for universe '%s'.", universe));
+              "Configuration operations are not currently supported for universe '%s'.",
+              universe instanceof UniverseScope.RoutedUniverse routed
+                  ? routed.atsControllerId()
+                  : "self"));
     }
 
     if (!Flags.instance().feConnectToConfigServer.getNonNull()) {

src/java/com/google/devtools/mobileharness/fe/v6/service/config/util/ConfigServiceCapabilityFactory.java

@@ -16,8 +16,19 @@
 
 package com.google.devtools.mobileharness.fe.v6.service.config.util;
 
+import com.google.devtools.mobileharness.fe.v6.service.util.UniverseScope;
+
 /** Factory for creating {@link ConfigServiceCapability} instances. */
 public interface ConfigServiceCapabilityFactory {
 
-  ConfigServiceCapability create(String universe);
+  ConfigServiceCapability create(UniverseScope universe);
+
+  /**
+   * @deprecated Use {@link #create(UniverseScope)} instead. TODO: Remove after all callers are
+   *     migrated to UniverseScope.
+   */
+  @Deprecated
+  default ConfigServiceCapability create(String universe) {
+    return create(UniverseScope.fromString(universe));
+  }
 }

src/javatests/com/google/devtools/mobileharness/fe/v6/service/config/util/ConfigServiceCapabilityTest.java

@@ -24,6 +24,7 @@
 import com.google.devtools.mobileharness.fe.v6.service.proto.config.DeviceConfigUiStatus;
 import com.google.devtools.mobileharness.fe.v6.service.proto.config.HostConfigUiStatus;
 import com.google.devtools.mobileharness.fe.v6.service.util.Environment;
+import com.google.devtools.mobileharness.fe.v6.service.util.UniverseScope;
 import com.google.devtools.mobileharness.shared.util.junit.rule.SetFlagsOss;
 import org.junit.Before;
 import org.junit.Rule;
@@ -42,6 +43,11 @@ public final class ConfigServiceCapabilityTest {
 
   @Mock private Environment environment;
 
+  private static final UniverseScope SELF_UNIVERSE = new UniverseScope.SelfUniverse();
+  private static final UniverseScope ROUTED_UNIVERSE = new UniverseScope.RoutedUniverse("other");
+  private static final UniverseScope UNSUPPORTED_UNIVERSE =
+      new UniverseScope.RoutedUniverse("unsupported_env");
+
   @Before
   public void setUp() {
     flags.setAllFlags(ImmutableMap.of("fe_connect_to_config_server", "true"));
@@ -51,49 +57,52 @@ public void setUp() {
   public void isConfigServiceAvailable_googleInternal_supported() {
     when(environment.isGoogleInternal()).thenReturn(true);
 
-    assertThat(new ConfigServiceCapability(environment, "google_1p").isConfigServiceAvailable())
+    assertThat(new ConfigServiceCapability(environment, SELF_UNIVERSE).isConfigServiceAvailable())
+        .isTrue();
+    assertThat(new ConfigServiceCapability(environment, SELF_UNIVERSE).isConfigServiceAvailable())
         .isTrue();
-    assertThat(new ConfigServiceCapability(environment, "").isConfigServiceAvailable()).isTrue();
   }
 
   @Test
   public void isConfigServiceAvailable_googleInternal_flagDisabled() {
     when(environment.isGoogleInternal()).thenReturn(true);
     flags.setAllFlags(ImmutableMap.of("fe_connect_to_config_server", "false"));
 
-    assertThat(new ConfigServiceCapability(environment, "google_1p").isConfigServiceAvailable())
+    assertThat(new ConfigServiceCapability(environment, SELF_UNIVERSE).isConfigServiceAvailable())
+        .isTrue();
+    assertThat(new ConfigServiceCapability(environment, SELF_UNIVERSE).isConfigServiceAvailable())
         .isTrue();
-    assertThat(new ConfigServiceCapability(environment, "").isConfigServiceAvailable()).isTrue();
   }
 
   @Test
   public void isConfigServiceAvailable_googleInternal_unsupported() {
     when(environment.isGoogleInternal()).thenReturn(true);
 
-    assertThat(new ConfigServiceCapability(environment, "other").isConfigServiceAvailable())
+    assertThat(new ConfigServiceCapability(environment, ROUTED_UNIVERSE).isConfigServiceAvailable())
         .isFalse();
   }
 
   @Test
   public void isConfigServiceAvailable_notGoogleInternal_supported() {
     when(environment.isGoogleInternal()).thenReturn(false);
 
-    assertThat(new ConfigServiceCapability(environment, "").isConfigServiceAvailable()).isTrue();
+    assertThat(new ConfigServiceCapability(environment, SELF_UNIVERSE).isConfigServiceAvailable())
+        .isTrue();
   }
 
   @Test
   public void isConfigServiceAvailable_notGoogleInternal_unsupported() {
     when(environment.isGoogleInternal()).thenReturn(false);
 
-    assertThat(new ConfigServiceCapability(environment, "other").isConfigServiceAvailable())
+    assertThat(new ConfigServiceCapability(environment, ROUTED_UNIVERSE).isConfigServiceAvailable())
         .isFalse();
   }
 
   @Test
   public void checkConfigServiceAvailability_googleInternal_unsupported() {
     when(environment.isGoogleInternal()).thenReturn(true);
     ConfigServiceCapability capability =
-        new ConfigServiceCapability(environment, "unsupported_env");
+        new ConfigServiceCapability(environment, UNSUPPORTED_UNIVERSE);
 
     UnsupportedOperationException exception =
         assertThrows(
@@ -109,7 +118,7 @@ public void checkConfigServiceAvailability_googleInternal_unsupported() {
   public void checkConfigServiceAvailability_notGoogleInternal_unsupported() {
     when(environment.isGoogleInternal()).thenReturn(false);
     ConfigServiceCapability capability =
-        new ConfigServiceCapability(environment, "unsupported_env");
+        new ConfigServiceCapability(environment, UNSUPPORTED_UNIVERSE);
 
     UnsupportedOperationException exception =
         assertThrows(
@@ -123,7 +132,7 @@ public void checkConfigServiceAvailability_notGoogleInternal_unsupported() {
   @Test
   public void checkConfigServiceAvailability_disabledByFlag() {
     flags.setAllFlags(ImmutableMap.of("fe_connect_to_config_server", "false"));
-    ConfigServiceCapability capability = new ConfigServiceCapability(environment, "");
+    ConfigServiceCapability capability = new ConfigServiceCapability(environment, SELF_UNIVERSE);
 
     UnsupportedOperationException exception =
         assertThrows(
@@ -137,7 +146,7 @@ public void checkConfigServiceAvailability_disabledByFlag() {
   @Test
   public void calculateHostUiStatus_ats() {
     when(environment.isAts()).thenReturn(true);
-    ConfigServiceCapability capability = new ConfigServiceCapability(environment, "");
+    ConfigServiceCapability capability = new ConfigServiceCapability(environment, SELF_UNIVERSE);
 
     HostConfigUiStatus status = capability.calculateHostUiStatus();
 
@@ -151,7 +160,7 @@ public void calculateHostUiStatus_ats() {
   @Test
   public void calculateHostUiStatus_notAts() {
     when(environment.isAts()).thenReturn(false);
-    ConfigServiceCapability capability = new ConfigServiceCapability(environment, "");
+    ConfigServiceCapability capability = new ConfigServiceCapability(environment, SELF_UNIVERSE);
 
     HostConfigUiStatus status = capability.calculateHostUiStatus();
 
@@ -166,7 +175,7 @@ public void calculateHostUiStatus_notAts() {
   @Test
   public void calculateDeviceUiStatus_ats() {
     when(environment.isAts()).thenReturn(true);
-    ConfigServiceCapability capability = new ConfigServiceCapability(environment, "");
+    ConfigServiceCapability capability = new ConfigServiceCapability(environment, SELF_UNIVERSE);
 
     DeviceConfigUiStatus status = capability.calculateDeviceUiStatus();
 
@@ -179,7 +188,7 @@ public void calculateDeviceUiStatus_ats() {
   @Test
   public void calculateDeviceUiStatus_notAts() {
     when(environment.isAts()).thenReturn(false);
-    ConfigServiceCapability capability = new ConfigServiceCapability(environment, "");
+    ConfigServiceCapability capability = new ConfigServiceCapability(environment, SELF_UNIVERSE);
 
     DeviceConfigUiStatus status = capability.calculateDeviceUiStatus();
 

src/javatests/com/google/devtools/mobileharness/fe/v6/service/shared/DeviceDataLoaderTest.java

@@ -21,6 +21,7 @@
 import static com.google.common.util.concurrent.Futures.immediateFuture;
 import static com.google.common.util.concurrent.MoreExecutors.newDirectExecutorService;
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.Mockito.when;
 
 import com.google.common.util.concurrent.ListenableFuture;
@@ -104,7 +105,7 @@ public final class DeviceDataLoaderTest {
   public void setUp() {
     Guice.createInjector(BoundFieldModule.of(this)).injectMembers(this);
 
-    when(configServiceCapabilityFactory.create(any())).thenReturn(configServiceCapability);
+    when(configServiceCapabilityFactory.create(anyString())).thenReturn(configServiceCapability);
     when(labInfoProvider.getLabInfoAsync(any(GetLabInfoRequest.class), any()))
         .thenReturn(immediateFuture(DEFAULT_LAB_INFO_RESPONSE));
     when(configurationProvider.getDeviceConfig(any(), any()))