From 756561fa132a24ad20fd4af1f1fad3a8319152ea Mon Sep 17 00:00:00 2001 From: Joel Thibault Date: Thu, 3 Oct 2024 16:24:56 -0400 Subject: [PATCH 1/4] don't bimap LeonardoMapper and LeonardoLabelHelper refactoring --- .../workbench/api/RuntimeController.java | 6 +- .../leonardo/LeonardoApiClientImpl.java | 5 +- .../leonardo/LeonardoLabelHelper.java | 12 ++++ .../utils/mappers/LeonardoMapper.java | 20 +++--- .../workbench/api/RuntimeControllerTest.java | 62 +++++++++++-------- 5 files changed, 62 insertions(+), 43 deletions(-) diff --git a/api/src/main/java/org/pmiops/workbench/api/RuntimeController.java b/api/src/main/java/org/pmiops/workbench/api/RuntimeController.java index ca4ac3e8543..514e09601e7 100644 --- a/api/src/main/java/org/pmiops/workbench/api/RuntimeController.java +++ b/api/src/main/java/org/pmiops/workbench/api/RuntimeController.java @@ -24,6 +24,7 @@ import org.pmiops.workbench.interactiveanalysis.InteractiveAnalysisService; import org.pmiops.workbench.leonardo.LeonardoApiClient; import org.pmiops.workbench.leonardo.LeonardoApiHelper; +import org.pmiops.workbench.leonardo.LeonardoLabelHelper; import org.pmiops.workbench.leonardo.PersistentDiskUtils; import org.pmiops.workbench.leonardo.model.LeonardoClusterError; import org.pmiops.workbench.leonardo.model.LeonardoGetRuntimeResponse; @@ -138,9 +139,8 @@ private Runtime getOverrideFromListRuntimes(String googleProject) { Map runtimeLabels = (Map) mostRecentRuntime.getLabels(); if (runtimeLabels != null - && LeonardoMapper.RUNTIME_CONFIGURATION_TYPE_ENUM_TO_STORAGE_MAP - .values() - .contains(runtimeLabels.get(LEONARDO_LABEL_AOU_CONFIG))) { + && LeonardoLabelHelper.isValidRuntimeConfigurationLabel( + runtimeLabels.get(LEONARDO_LABEL_AOU_CONFIG))) { try { Runtime runtime = leonardoMapper.toApiRuntime(mostRecentRuntime); if (!RuntimeStatus.DELETED.equals(runtime.getStatus())) { diff --git a/api/src/main/java/org/pmiops/workbench/leonardo/LeonardoApiClientImpl.java b/api/src/main/java/org/pmiops/workbench/leonardo/LeonardoApiClientImpl.java index fd1c3eb0860..ea0f341ec4c 100644 --- a/api/src/main/java/org/pmiops/workbench/leonardo/LeonardoApiClientImpl.java +++ b/api/src/main/java/org/pmiops/workbench/leonardo/LeonardoApiClientImpl.java @@ -4,6 +4,7 @@ import static org.pmiops.workbench.leonardo.LeonardoCustomEnvVarUtils.OWNER_EMAIL_ENV_KEY; import static org.pmiops.workbench.leonardo.LeonardoCustomEnvVarUtils.WORKSPACE_NAME_ENV_KEY; import static org.pmiops.workbench.leonardo.LeonardoLabelHelper.LEONARDO_DISK_LABEL_KEYS; +import static org.pmiops.workbench.leonardo.LeonardoLabelHelper.LEONARDO_LABEL_AOU_CONFIG; import static org.pmiops.workbench.leonardo.LeonardoLabelHelper.appTypeToLabelValue; import static org.pmiops.workbench.leonardo.LeonardoLabelHelper.upsertLeonardoLabel; @@ -303,9 +304,7 @@ private Map buildRuntimeConfigurationLabels( RuntimeConfigurationType runtimeConfigurationType) { if (runtimeConfigurationType != null) { return Collections.singletonMap( - LeonardoLabelHelper.LEONARDO_LABEL_AOU_CONFIG, - LeonardoMapper.RUNTIME_CONFIGURATION_TYPE_ENUM_TO_STORAGE_MAP.get( - runtimeConfigurationType)); + LEONARDO_LABEL_AOU_CONFIG, leonardoMapper.toConfigurationLabel(runtimeConfigurationType)); } else { return new HashMap<>(); } diff --git a/api/src/main/java/org/pmiops/workbench/leonardo/LeonardoLabelHelper.java b/api/src/main/java/org/pmiops/workbench/leonardo/LeonardoLabelHelper.java index 5624652fdc6..24b938ac6f0 100644 --- a/api/src/main/java/org/pmiops/workbench/leonardo/LeonardoLabelHelper.java +++ b/api/src/main/java/org/pmiops/workbench/leonardo/LeonardoLabelHelper.java @@ -4,6 +4,7 @@ import java.util.HashMap; import java.util.Map; import java.util.Optional; +import java.util.Set; import org.pmiops.workbench.model.AppType; /** Helper class for setting Leonardo labels. */ @@ -26,6 +27,17 @@ private LeonardoLabelHelper() {} public static final String LEONARDO_LABEL_WORKSPACE_NAMESPACE = "saturnWorkspaceNamespace"; public static final String LEONARDO_LABEL_WORKSPACE_NAME = "saturnWorkspaceName"; + // Important: keep these string constants in sync with LeonardoMapper + // toConfigurationType() and toConfigurationLabel() + + public static String USER_OVERRIDE = "user-override"; + public static String GENERAL_ANALYSIS = "preset-general-analysis"; + public static String HAIL_GENOMIC_ANALYSIS = "preset-hail-genomic-analysis"; + + public static boolean isValidRuntimeConfigurationLabel(String s) { + return Set.of(USER_OVERRIDE, GENERAL_ANALYSIS, HAIL_GENOMIC_ANALYSIS).contains(s); + } + public static String appTypeToLabelValue(AppType appType) { return appType.toString().toLowerCase(); } diff --git a/api/src/main/java/org/pmiops/workbench/utils/mappers/LeonardoMapper.java b/api/src/main/java/org/pmiops/workbench/utils/mappers/LeonardoMapper.java index c1244c966c9..c115820bf13 100644 --- a/api/src/main/java/org/pmiops/workbench/utils/mappers/LeonardoMapper.java +++ b/api/src/main/java/org/pmiops/workbench/utils/mappers/LeonardoMapper.java @@ -1,7 +1,5 @@ package org.pmiops.workbench.utils.mappers; -import com.google.common.collect.BiMap; -import com.google.common.collect.ImmutableBiMap; import com.google.gson.Gson; import jakarta.annotation.Nullable; import java.util.List; @@ -57,11 +55,15 @@ @Mapper(config = MapStructConfig.class) public interface LeonardoMapper { - BiMap RUNTIME_CONFIGURATION_TYPE_ENUM_TO_STORAGE_MAP = - ImmutableBiMap.of( - RuntimeConfigurationType.USEROVERRIDE, "user-override", - RuntimeConfigurationType.GENERALANALYSIS, "preset-general-analysis", - RuntimeConfigurationType.HAILGENOMICANALYSIS, "preset-hail-genomic-analysis"); + @ValueMapping(source = "user-override", target = "USEROVERRIDE") + @ValueMapping(source = "preset-general-analysis", target = "GENERALANALYSIS") + @ValueMapping(source = "preset-hail-genomic-analysis", target = "HAILGENOMICANALYSIS") + RuntimeConfigurationType toConfigurationType(String runtimeConfigurationLabel); + + @ValueMapping(source = "USEROVERRIDE", target = "user-override") + @ValueMapping(source = "GENERALANALYSIS", target = "preset-general-analysis") + @ValueMapping(source = "HAILGENOMICANALYSIS", target = "preset-hail-genomic-analysis") + String toConfigurationLabel(RuntimeConfigurationType runtimeConfigurationType); DataprocConfig toDataprocConfig(LeonardoMachineConfig leonardoMachineConfig); @@ -267,9 +269,7 @@ default void mapRuntimeLabels(Runtime runtime, Object runtimeLabelsObj) { runtime.setConfigurationType(RuntimeConfigurationType.HAILGENOMICANALYSIS); } else { runtime.setConfigurationType( - RUNTIME_CONFIGURATION_TYPE_ENUM_TO_STORAGE_MAP - .inverse() - .get(runtimeLabels.get(LeonardoLabelHelper.LEONARDO_LABEL_AOU_CONFIG))); + toConfigurationType(runtimeLabels.get(LeonardoLabelHelper.LEONARDO_LABEL_AOU_CONFIG))); } } diff --git a/api/src/test/java/org/pmiops/workbench/api/RuntimeControllerTest.java b/api/src/test/java/org/pmiops/workbench/api/RuntimeControllerTest.java index 0aee0a50523..d4fe23fc0ae 100644 --- a/api/src/test/java/org/pmiops/workbench/api/RuntimeControllerTest.java +++ b/api/src/test/java/org/pmiops/workbench/api/RuntimeControllerTest.java @@ -10,11 +10,14 @@ import static org.mockito.Mockito.never; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; +import static org.pmiops.workbench.leonardo.LeonardoLabelHelper.GENERAL_ANALYSIS; +import static org.pmiops.workbench.leonardo.LeonardoLabelHelper.HAIL_GENOMIC_ANALYSIS; import static org.pmiops.workbench.leonardo.LeonardoLabelHelper.LEONARDO_LABEL_AOU_CONFIG; import static org.pmiops.workbench.leonardo.LeonardoLabelHelper.LEONARDO_LABEL_IS_RUNTIME; import static org.pmiops.workbench.leonardo.LeonardoLabelHelper.LEONARDO_LABEL_IS_RUNTIME_TRUE; import static org.pmiops.workbench.leonardo.LeonardoLabelHelper.LEONARDO_LABEL_WORKSPACE_NAME; import static org.pmiops.workbench.leonardo.LeonardoLabelHelper.LEONARDO_LABEL_WORKSPACE_NAMESPACE; +import static org.pmiops.workbench.leonardo.LeonardoLabelHelper.USER_OVERRIDE; import static org.pmiops.workbench.utils.TestMockFactory.createControlledTier; import com.google.cloud.Date; @@ -473,7 +476,7 @@ public void testGetRuntime_noLabel() throws ApiException { @Test public void testGetRuntime_defaultLabel_hail() throws ApiException { - testLeoRuntime.setLabels(ImmutableMap.of("all-of-us-config", "preset-hail-genomic-analysis")); + testLeoRuntime.setLabels(ImmutableMap.of(LEONARDO_LABEL_AOU_CONFIG, HAIL_GENOMIC_ANALYSIS)); when(mockUserRuntimesApi.getRuntime(GOOGLE_PROJECT_ID, getRuntimeName())) .thenReturn(testLeoRuntime); @@ -484,7 +487,10 @@ public void testGetRuntime_defaultLabel_hail() throws ApiException { @Test public void testGetRuntime_defaultLabel_generalAnalysis() throws ApiException { - testLeoRuntime.setLabels(ImmutableMap.of("all-of-us-config", "preset-general-analysis")); + testLeoRuntime.setLabels(ImmutableMap.of(LEONARDO_LABEL_AOU_CONFIG, GENERAL_ANALYSIS)); + + when(mockUserRuntimesApi.getRuntime(GOOGLE_PROJECT_ID, getRuntimeName())) + .thenReturn(testLeoRuntime); when(mockUserRuntimesApi.getRuntime(GOOGLE_PROJECT_ID, getRuntimeName())) .thenReturn(testLeoRuntime); @@ -495,7 +501,7 @@ public void testGetRuntime_defaultLabel_generalAnalysis() throws ApiException { @Test public void testGetRuntime_overrideLabel() throws ApiException { - testLeoRuntime.setLabels(ImmutableMap.of("all-of-us-config", "user-override")); + testLeoRuntime.setLabels(ImmutableMap.of(LEONARDO_LABEL_AOU_CONFIG, USER_OVERRIDE)); when(mockUserRuntimesApi.getRuntime(GOOGLE_PROJECT_ID, getRuntimeName())) .thenReturn(testLeoRuntime); @@ -531,7 +537,7 @@ public void testGetRuntime_fromListRuntimes() throws ApiException { .runtimeName("expected-runtime") .status(LeonardoRuntimeStatus.CREATING) .auditInfo(new LeonardoAuditInfo().createdDate(timestamp)) - .labels(ImmutableMap.of("all-of-us-config", "user-override")))); + .labels(ImmutableMap.of(LEONARDO_LABEL_AOU_CONFIG, USER_OVERRIDE)))); Runtime runtime = runtimeController.getRuntime(WORKSPACE_NS).getBody(); @@ -552,7 +558,7 @@ public void testGetRuntime_fromListRuntimes_invalidRuntime() throws ApiException ImmutableList.of( new LeonardoListRuntimeResponse() .runtimeConfig(dataprocConfigObj) - .labels(ImmutableMap.of("all-of-us-config", "user-override")))); + .labels(ImmutableMap.of(LEONARDO_LABEL_AOU_CONFIG, USER_OVERRIDE)))); assertThrows(NotFoundException.class, () -> runtimeController.getRuntime(WORKSPACE_NS)); } @@ -569,7 +575,7 @@ public void testGetRuntime_fromListRuntimes_gceConfig() throws ApiException { new LeonardoListRuntimeResponse() .runtimeConfig(gceConfigObj) .auditInfo(new LeonardoAuditInfo().createdDate(timestamp)) - .labels(ImmutableMap.of("all-of-us-config", "user-override")))); + .labels(ImmutableMap.of(LEONARDO_LABEL_AOU_CONFIG, USER_OVERRIDE)))); Runtime runtime = runtimeController.getRuntime(WORKSPACE_NS).getBody(); @@ -599,7 +605,7 @@ public void testGetRuntime_fromListRuntimes_dataprocConfig() throws ApiException new LeonardoListRuntimeResponse() .runtimeConfig(dataProcConfigObj) .auditInfo(new LeonardoAuditInfo().createdDate(timestamp)) - .labels(ImmutableMap.of("all-of-us-config", "user-override")))); + .labels(ImmutableMap.of(LEONARDO_LABEL_AOU_CONFIG, USER_OVERRIDE)))); Runtime runtime = runtimeController.getRuntime(WORKSPACE_NS).getBody(); @@ -626,11 +632,11 @@ public void testGetRuntime_fromListRuntimes_checkMostRecent() throws ApiExceptio new LeonardoListRuntimeResponse() .runtimeName("expected-runtime") .auditInfo(new LeonardoAuditInfo().createdDate(newerTimestamp)) - .labels(ImmutableMap.of("all-of-us-config", "user-override")), + .labels(ImmutableMap.of(LEONARDO_LABEL_AOU_CONFIG, USER_OVERRIDE)), new LeonardoListRuntimeResponse() .runtimeName("default-runtime") .auditInfo(new LeonardoAuditInfo().createdDate(olderTimestamp)) - .labels(ImmutableMap.of("all-of-us-config", "default")))); + .labels(ImmutableMap.of(LEONARDO_LABEL_AOU_CONFIG, "default")))); assertThat(runtimeController.getRuntime(WORKSPACE_NS).getBody().getRuntimeName()) .isEqualTo("expected-runtime"); @@ -648,10 +654,10 @@ public void testGetRuntime_fromListRuntimes_checkMostRecent_nullAuditInfo() thro new LeonardoListRuntimeResponse() .runtimeName("expected-runtime") .auditInfo(new LeonardoAuditInfo().createdDate(newerTimestamp)) - .labels(ImmutableMap.of("all-of-us-config", "user-override")), + .labels(ImmutableMap.of(LEONARDO_LABEL_AOU_CONFIG, USER_OVERRIDE)), new LeonardoListRuntimeResponse() .runtimeName("default-runtime") - .labels(ImmutableMap.of("all-of-us-config", "default")))); + .labels(ImmutableMap.of(LEONARDO_LABEL_AOU_CONFIG, "default")))); assertThat(runtimeController.getRuntime(WORKSPACE_NS).getBody().getRuntimeName()) .isEqualTo("expected-runtime"); @@ -669,11 +675,11 @@ public void testGetRuntime_fromListRuntimes_checkMostRecent_nullTimestamp() thro new LeonardoListRuntimeResponse() .runtimeName("expected-runtime") .auditInfo(new LeonardoAuditInfo().createdDate(newerTimestamp)) - .labels(ImmutableMap.of("all-of-us-config", "user-override")), + .labels(ImmutableMap.of(LEONARDO_LABEL_AOU_CONFIG, USER_OVERRIDE)), new LeonardoListRuntimeResponse() .runtimeName("default-runtime") .auditInfo(new LeonardoAuditInfo().createdDate(null)) - .labels(ImmutableMap.of("all-of-us-config", "default")))); + .labels(ImmutableMap.of(LEONARDO_LABEL_AOU_CONFIG, "default")))); assertThat(runtimeController.getRuntime(WORKSPACE_NS).getBody().getRuntimeName()) .isEqualTo("expected-runtime"); @@ -691,11 +697,11 @@ public void testGetRuntime_fromListRuntimes_checkMostRecent_emptyTimestamp() thr new LeonardoListRuntimeResponse() .runtimeName("expected-runtime") .auditInfo(new LeonardoAuditInfo().createdDate(newerTimestamp)) - .labels(ImmutableMap.of("all-of-us-config", "user-override")), + .labels(ImmutableMap.of(LEONARDO_LABEL_AOU_CONFIG, USER_OVERRIDE)), new LeonardoListRuntimeResponse() .runtimeName("default-runtime") .auditInfo(new LeonardoAuditInfo().createdDate("")) - .labels(ImmutableMap.of("all-of-us-config", "default")))); + .labels(ImmutableMap.of(LEONARDO_LABEL_AOU_CONFIG, "default")))); assertThat(runtimeController.getRuntime(WORKSPACE_NS).getBody().getRuntimeName()) .isEqualTo("expected-runtime"); @@ -714,11 +720,11 @@ public void testGetRuntime_fromListRuntime_mostRecentIsDefaultLabel() throws Api new LeonardoListRuntimeResponse() .runtimeName("override-runtime") .auditInfo(new LeonardoAuditInfo().createdDate(olderTimestamp)) - .labels(ImmutableMap.of("all-of-us-config", "user-override")), + .labels(ImmutableMap.of(LEONARDO_LABEL_AOU_CONFIG, USER_OVERRIDE)), new LeonardoListRuntimeResponse() .runtimeName("default-runtime") .auditInfo(new LeonardoAuditInfo().createdDate(newerTimestamp)) - .labels(ImmutableMap.of("all-of-us-config", "default")))); + .labels(ImmutableMap.of(LEONARDO_LABEL_AOU_CONFIG, "default")))); assertThrows(NotFoundException.class, () -> runtimeController.getRuntime(WORKSPACE_NS)); } @@ -739,7 +745,7 @@ public void testGetRuntime_fromListRuntime_mostRecentHasNoLabel() throws ApiExce new LeonardoListRuntimeResponse() .runtimeName("default-runtime") .auditInfo(new LeonardoAuditInfo().createdDate(olderTimestamp)) - .labels(ImmutableMap.of("all-of-us-config", "default")))); + .labels(ImmutableMap.of(LEONARDO_LABEL_AOU_CONFIG, "default")))); assertThrows(NotFoundException.class, () -> runtimeController.getRuntime(WORKSPACE_NS)); } @@ -756,7 +762,7 @@ public void testGetRuntime_fromListRuntime_returnPresets() throws ApiException { new LeonardoListRuntimeResponse() .runtimeName("preset-runtime") .auditInfo(new LeonardoAuditInfo().createdDate(timestamp)) - .labels(ImmutableMap.of("all-of-us-config", "preset-general-analysis")))); + .labels(ImmutableMap.of(LEONARDO_LABEL_AOU_CONFIG, GENERAL_ANALYSIS)))); assertThat(runtimeController.getRuntime(WORKSPACE_NS).getBody().getRuntimeName()) .isEqualTo("preset-runtime"); @@ -1047,8 +1053,9 @@ public void testCreateRuntime_defaultLabel_hail() throws ApiException { eq(GOOGLE_PROJECT_ID), eq(getRuntimeName()), createRuntimeRequestCaptor.capture()); LeonardoCreateRuntimeRequest createRuntimeRequest = createRuntimeRequestCaptor.getValue(); - assertThat(((Map) createRuntimeRequest.getLabels()).get("all-of-us-config")) - .isEqualTo("preset-hail-genomic-analysis"); + assertThat( + ((Map) createRuntimeRequest.getLabels()).get(LEONARDO_LABEL_AOU_CONFIG)) + .isEqualTo(HAIL_GENOMIC_ANALYSIS); } @Test @@ -1067,8 +1074,9 @@ public void testCreateRuntime_defaultLabel_generalAnalysis() throws ApiException eq(GOOGLE_PROJECT_ID), eq(getRuntimeName()), createRuntimeRequestCaptor.capture()); LeonardoCreateRuntimeRequest createRuntimeRequest = createRuntimeRequestCaptor.getValue(); - assertThat(((Map) createRuntimeRequest.getLabels()).get("all-of-us-config")) - .isEqualTo("preset-general-analysis"); + assertThat( + ((Map) createRuntimeRequest.getLabels()).get(LEONARDO_LABEL_AOU_CONFIG)) + .isEqualTo(GENERAL_ANALYSIS); } @Test @@ -1087,8 +1095,9 @@ public void testCreateRuntime_overrideLabel() throws ApiException { eq(GOOGLE_PROJECT_ID), eq(getRuntimeName()), createRuntimeRequestCaptor.capture()); LeonardoCreateRuntimeRequest createRuntimeRequest = createRuntimeRequestCaptor.getValue(); - assertThat(((Map) createRuntimeRequest.getLabels()).get("all-of-us-config")) - .isEqualTo("user-override"); + assertThat( + ((Map) createRuntimeRequest.getLabels()).get(LEONARDO_LABEL_AOU_CONFIG)) + .isEqualTo(USER_OVERRIDE); } @Test @@ -1308,8 +1317,7 @@ public void testUpdateRuntime() throws ApiException { .isEqualTo( Collections.singletonMap( LEONARDO_LABEL_AOU_CONFIG, - LeonardoMapper.RUNTIME_CONFIGURATION_TYPE_ENUM_TO_STORAGE_MAP.get( - RuntimeConfigurationType.USEROVERRIDE))); + leonardoMapper.toConfigurationLabel(RuntimeConfigurationType.USEROVERRIDE))); } @Test From ec3081d746fc43e4767a126b57dab5fd14c64758 Mon Sep 17 00:00:00 2001 From: Joel Thibault Date: Thu, 3 Oct 2024 16:54:11 -0400 Subject: [PATCH 2/4] oops --- .../java/org/pmiops/workbench/api/RuntimeControllerTest.java | 3 --- 1 file changed, 3 deletions(-) diff --git a/api/src/test/java/org/pmiops/workbench/api/RuntimeControllerTest.java b/api/src/test/java/org/pmiops/workbench/api/RuntimeControllerTest.java index d4fe23fc0ae..f0b3f6a3956 100644 --- a/api/src/test/java/org/pmiops/workbench/api/RuntimeControllerTest.java +++ b/api/src/test/java/org/pmiops/workbench/api/RuntimeControllerTest.java @@ -489,9 +489,6 @@ public void testGetRuntime_defaultLabel_hail() throws ApiException { public void testGetRuntime_defaultLabel_generalAnalysis() throws ApiException { testLeoRuntime.setLabels(ImmutableMap.of(LEONARDO_LABEL_AOU_CONFIG, GENERAL_ANALYSIS)); - when(mockUserRuntimesApi.getRuntime(GOOGLE_PROJECT_ID, getRuntimeName())) - .thenReturn(testLeoRuntime); - when(mockUserRuntimesApi.getRuntime(GOOGLE_PROJECT_ID, getRuntimeName())) .thenReturn(testLeoRuntime); From bd477ebe11e7cb8a84e5d45d3eb3225b97ed03ee Mon Sep 17 00:00:00 2001 From: Joel Thibault Date: Thu, 3 Oct 2024 17:06:14 -0400 Subject: [PATCH 3/4] better --- .../workbench/api/RuntimeController.java | 7 ++---- .../leonardo/LeonardoLabelHelper.java | 25 ++++++++++++++++++- .../workbench/api/RuntimeControllerTest.java | 10 +++----- 3 files changed, 30 insertions(+), 12 deletions(-) diff --git a/api/src/main/java/org/pmiops/workbench/api/RuntimeController.java b/api/src/main/java/org/pmiops/workbench/api/RuntimeController.java index 514e09601e7..48e8f29c904 100644 --- a/api/src/main/java/org/pmiops/workbench/api/RuntimeController.java +++ b/api/src/main/java/org/pmiops/workbench/api/RuntimeController.java @@ -1,10 +1,10 @@ package org.pmiops.workbench.api; -import static org.pmiops.workbench.leonardo.LeonardoLabelHelper.LEONARDO_LABEL_AOU_CONFIG; import static org.pmiops.workbench.leonardo.LeonardoLabelHelper.LEONARDO_LABEL_IS_RUNTIME; import static org.pmiops.workbench.leonardo.LeonardoLabelHelper.LEONARDO_LABEL_IS_RUNTIME_TRUE; import static org.pmiops.workbench.leonardo.LeonardoLabelHelper.LEONARDO_LABEL_WORKSPACE_NAME; import static org.pmiops.workbench.leonardo.LeonardoLabelHelper.LEONARDO_LABEL_WORKSPACE_NAMESPACE; +import static org.pmiops.workbench.leonardo.LeonardoLabelHelper.hasValidRuntimeConfigurationLabel; import static org.pmiops.workbench.leonardo.LeonardoLabelHelper.upsertLeonardoLabel; import com.google.common.base.Strings; @@ -24,7 +24,6 @@ import org.pmiops.workbench.interactiveanalysis.InteractiveAnalysisService; import org.pmiops.workbench.leonardo.LeonardoApiClient; import org.pmiops.workbench.leonardo.LeonardoApiHelper; -import org.pmiops.workbench.leonardo.LeonardoLabelHelper; import org.pmiops.workbench.leonardo.PersistentDiskUtils; import org.pmiops.workbench.leonardo.model.LeonardoClusterError; import org.pmiops.workbench.leonardo.model.LeonardoGetRuntimeResponse; @@ -138,9 +137,7 @@ private Runtime getOverrideFromListRuntimes(String googleProject) { @SuppressWarnings("unchecked") Map runtimeLabels = (Map) mostRecentRuntime.getLabels(); - if (runtimeLabels != null - && LeonardoLabelHelper.isValidRuntimeConfigurationLabel( - runtimeLabels.get(LEONARDO_LABEL_AOU_CONFIG))) { + if (hasValidRuntimeConfigurationLabel(runtimeLabels)) { try { Runtime runtime = leonardoMapper.toApiRuntime(mostRecentRuntime); if (!RuntimeStatus.DELETED.equals(runtime.getStatus())) { diff --git a/api/src/main/java/org/pmiops/workbench/leonardo/LeonardoLabelHelper.java b/api/src/main/java/org/pmiops/workbench/leonardo/LeonardoLabelHelper.java index 24b938ac6f0..139c1dd6a52 100644 --- a/api/src/main/java/org/pmiops/workbench/leonardo/LeonardoLabelHelper.java +++ b/api/src/main/java/org/pmiops/workbench/leonardo/LeonardoLabelHelper.java @@ -34,7 +34,30 @@ private LeonardoLabelHelper() {} public static String GENERAL_ANALYSIS = "preset-general-analysis"; public static String HAIL_GENOMIC_ANALYSIS = "preset-hail-genomic-analysis"; - public static boolean isValidRuntimeConfigurationLabel(String s) { + @Nullable + public static String getRuntimeConfigurationLabel(@Nullable Map labels) { + if (labels == null) { + return null; + } + return labels.get(LEONARDO_LABEL_AOU_CONFIG); + } + + @Nullable + @SuppressWarnings("unchecked") + public static String getRuntimeConfigurationLabel(@Nullable Object labels) { + if (labels == null) { + return null; + } + + return getRuntimeConfigurationLabel((Map) labels); + } + + public static boolean hasValidRuntimeConfigurationLabel(@Nullable Map labels) { + String s = getRuntimeConfigurationLabel(labels); + if (s == null) { + return false; + } + return Set.of(USER_OVERRIDE, GENERAL_ANALYSIS, HAIL_GENOMIC_ANALYSIS).contains(s); } diff --git a/api/src/test/java/org/pmiops/workbench/api/RuntimeControllerTest.java b/api/src/test/java/org/pmiops/workbench/api/RuntimeControllerTest.java index f0b3f6a3956..5fea467451b 100644 --- a/api/src/test/java/org/pmiops/workbench/api/RuntimeControllerTest.java +++ b/api/src/test/java/org/pmiops/workbench/api/RuntimeControllerTest.java @@ -18,6 +18,7 @@ import static org.pmiops.workbench.leonardo.LeonardoLabelHelper.LEONARDO_LABEL_WORKSPACE_NAME; import static org.pmiops.workbench.leonardo.LeonardoLabelHelper.LEONARDO_LABEL_WORKSPACE_NAMESPACE; import static org.pmiops.workbench.leonardo.LeonardoLabelHelper.USER_OVERRIDE; +import static org.pmiops.workbench.leonardo.LeonardoLabelHelper.getRuntimeConfigurationLabel; import static org.pmiops.workbench.utils.TestMockFactory.createControlledTier; import com.google.cloud.Date; @@ -1050,8 +1051,7 @@ public void testCreateRuntime_defaultLabel_hail() throws ApiException { eq(GOOGLE_PROJECT_ID), eq(getRuntimeName()), createRuntimeRequestCaptor.capture()); LeonardoCreateRuntimeRequest createRuntimeRequest = createRuntimeRequestCaptor.getValue(); - assertThat( - ((Map) createRuntimeRequest.getLabels()).get(LEONARDO_LABEL_AOU_CONFIG)) + assertThat(getRuntimeConfigurationLabel(createRuntimeRequest.getLabels())) .isEqualTo(HAIL_GENOMIC_ANALYSIS); } @@ -1071,8 +1071,7 @@ public void testCreateRuntime_defaultLabel_generalAnalysis() throws ApiException eq(GOOGLE_PROJECT_ID), eq(getRuntimeName()), createRuntimeRequestCaptor.capture()); LeonardoCreateRuntimeRequest createRuntimeRequest = createRuntimeRequestCaptor.getValue(); - assertThat( - ((Map) createRuntimeRequest.getLabels()).get(LEONARDO_LABEL_AOU_CONFIG)) + assertThat(getRuntimeConfigurationLabel(createRuntimeRequest.getLabels())) .isEqualTo(GENERAL_ANALYSIS); } @@ -1092,8 +1091,7 @@ public void testCreateRuntime_overrideLabel() throws ApiException { eq(GOOGLE_PROJECT_ID), eq(getRuntimeName()), createRuntimeRequestCaptor.capture()); LeonardoCreateRuntimeRequest createRuntimeRequest = createRuntimeRequestCaptor.getValue(); - assertThat( - ((Map) createRuntimeRequest.getLabels()).get(LEONARDO_LABEL_AOU_CONFIG)) + assertThat(getRuntimeConfigurationLabel(createRuntimeRequest.getLabels())) .isEqualTo(USER_OVERRIDE); } From 02f0b4866a49633c55a4aed7ae89640ada4fce5b Mon Sep 17 00:00:00 2001 From: Joel Thibault Date: Thu, 3 Oct 2024 17:15:01 -0400 Subject: [PATCH 4/4] better --- .../workbench/api/RuntimeController.java | 6 +----- .../leonardo/LeonardoLabelHelper.java | 2 +- .../utils/mappers/LeonardoMapper.java | 19 +++++++++---------- 3 files changed, 11 insertions(+), 16 deletions(-) diff --git a/api/src/main/java/org/pmiops/workbench/api/RuntimeController.java b/api/src/main/java/org/pmiops/workbench/api/RuntimeController.java index 48e8f29c904..34050f88af9 100644 --- a/api/src/main/java/org/pmiops/workbench/api/RuntimeController.java +++ b/api/src/main/java/org/pmiops/workbench/api/RuntimeController.java @@ -11,7 +11,6 @@ import jakarta.annotation.Nullable; import jakarta.inject.Provider; import java.util.List; -import java.util.Map; import java.util.Objects; import java.util.Optional; import java.util.logging.Logger; @@ -134,10 +133,7 @@ private Runtime getOverrideFromListRuntimes(String googleProject) { LeonardoListRuntimeResponse mostRecentRuntime = mostRecentRuntimeMaybe.orElseThrow(NotFoundException::new); - @SuppressWarnings("unchecked") - Map runtimeLabels = (Map) mostRecentRuntime.getLabels(); - - if (hasValidRuntimeConfigurationLabel(runtimeLabels)) { + if (hasValidRuntimeConfigurationLabel(mostRecentRuntime.getLabels())) { try { Runtime runtime = leonardoMapper.toApiRuntime(mostRecentRuntime); if (!RuntimeStatus.DELETED.equals(runtime.getStatus())) { diff --git a/api/src/main/java/org/pmiops/workbench/leonardo/LeonardoLabelHelper.java b/api/src/main/java/org/pmiops/workbench/leonardo/LeonardoLabelHelper.java index 139c1dd6a52..ad998bde320 100644 --- a/api/src/main/java/org/pmiops/workbench/leonardo/LeonardoLabelHelper.java +++ b/api/src/main/java/org/pmiops/workbench/leonardo/LeonardoLabelHelper.java @@ -52,7 +52,7 @@ public static String getRuntimeConfigurationLabel(@Nullable Object labels) { return getRuntimeConfigurationLabel((Map) labels); } - public static boolean hasValidRuntimeConfigurationLabel(@Nullable Map labels) { + public static boolean hasValidRuntimeConfigurationLabel(@Nullable Object labels) { String s = getRuntimeConfigurationLabel(labels); if (s == null) { return false; diff --git a/api/src/main/java/org/pmiops/workbench/utils/mappers/LeonardoMapper.java b/api/src/main/java/org/pmiops/workbench/utils/mappers/LeonardoMapper.java index c115820bf13..eb615e0436d 100644 --- a/api/src/main/java/org/pmiops/workbench/utils/mappers/LeonardoMapper.java +++ b/api/src/main/java/org/pmiops/workbench/utils/mappers/LeonardoMapper.java @@ -1,9 +1,11 @@ package org.pmiops.workbench.utils.mappers; +import static org.pmiops.workbench.leonardo.LeonardoLabelHelper.getRuntimeConfigurationLabel; +import static org.pmiops.workbench.leonardo.LeonardoLabelHelper.hasValidRuntimeConfigurationLabel; + import com.google.gson.Gson; import jakarta.annotation.Nullable; import java.util.List; -import java.util.Map; import java.util.Optional; import org.mapstruct.AfterMapping; import org.mapstruct.Mapper; @@ -260,16 +262,13 @@ default void setAppType(UserAppEnvironment appEnvironment, @Nullable Object appL } default void mapRuntimeLabels(Runtime runtime, Object runtimeLabelsObj) { - @SuppressWarnings("unchecked") - final Map runtimeLabels = (Map) runtimeLabelsObj; - if (runtimeLabels == null - || runtimeLabels.get(LeonardoLabelHelper.LEONARDO_LABEL_AOU_CONFIG) == null) { - // If there's no label, fall back onto the old behavior where every Runtime was created with a - // default Dataproc config - runtime.setConfigurationType(RuntimeConfigurationType.HAILGENOMICANALYSIS); - } else { + if (hasValidRuntimeConfigurationLabel(runtimeLabelsObj)) { runtime.setConfigurationType( - toConfigurationType(runtimeLabels.get(LeonardoLabelHelper.LEONARDO_LABEL_AOU_CONFIG))); + toConfigurationType(getRuntimeConfigurationLabel(runtimeLabelsObj))); + } else { + // If there's no valid label, fall back onto the old behavior where every Runtime was created + // with a default Dataproc config + runtime.setConfigurationType(RuntimeConfigurationType.HAILGENOMICANALYSIS); } }