Fix unit tests

DataBiosphere · Sep 10, 2024 · 362ceac · 362ceac
1 parent 36e62e1
commit 362ceac
Show file tree

Hide file tree

Showing 8 changed files with 72 additions and 491 deletions.
diff --git a/.../scala/org/broadinstitute/dsde/workbench/leonardo/http/service/RuntimeServiceInterp.scala b/.../scala/org/broadinstitute/dsde/workbench/leonardo/http/service/RuntimeServiceInterp.scala
@@ -179,7 +179,7 @@ class RuntimeServiceInterp[F[_]: Parallel](
               }
 
             runtime = convertToRuntime(
-              userInfo,
+              userEmail,
               petSA,
               CloudContext.Gcp(googleProject),
               runtimeName,
@@ -1068,7 +1068,7 @@ object RuntimeServiceInterp {
     }
 
   private[service] def convertToRuntime(
-    userInfo: UserInfo,
+    userEmail: WorkbenchEmail,
     serviceAccountInfo: WorkbenchEmail,
     cloudContext: CloudContext,
     runtimeName: RuntimeName,
@@ -1084,7 +1084,7 @@ object RuntimeServiceInterp {
       runtimeName,
       Some(GoogleProject(cloudContext.asString)),
       cloudContext,
-      userInfo.userEmail,
+      userEmail,
       Some(serviceAccountInfo),
       req.userScriptUri,
       req.startUserScriptUri,
@@ -1123,7 +1123,7 @@ object RuntimeServiceInterp {
       cloudContext = cloudContext,
       serviceAccount = serviceAccountInfo,
       asyncRuntimeFields = None,
-      auditInfo = AuditInfo(userInfo.userEmail, now, None, now),
+      auditInfo = AuditInfo(userEmail, now, None, now),
       kernelFoundBusyDate = None,
       proxyUrl = Runtime.getProxyUrl(config.proxyUrlBase, cloudContext, runtimeName, clusterImages, None, allLabels),
       status = RuntimeStatus.PreCreating,

diff --git a/...cala/org/broadinstitute/dsde/workbench/leonardo/http/service/RuntimeV2ServiceInterp.scala b/...cala/org/broadinstitute/dsde/workbench/leonardo/http/service/RuntimeV2ServiceInterp.scala
@@ -180,7 +180,7 @@ class RuntimeV2ServiceInterp[F[_]: Parallel](
                   samResource <- F.delay(PersistentDiskSamResourceId(UUID.randomUUID().toString))
                   pd <- F.fromEither(
                     convertToDisk(
-                      userInfo,
+                      userEmail,
                       cloudContext,
                       DiskName(req.azureDiskConfig.name.value),
                       samResource,
@@ -206,7 +206,7 @@ class RuntimeV2ServiceInterp[F[_]: Parallel](
               workspaceId,
               runtimeName,
               cloudContext,
-              userInfo,
+              userEmail,
               req,
               samResource,
               Set(runtimeImage, listenerImage, welderImage),
@@ -567,7 +567,7 @@ class RuntimeV2ServiceInterp[F[_]: Parallel](
     } yield runtimes.toVector
 
   private[service] def convertToDisk(
-    userInfo: UserInfo,
+    userEmail: WorkbenchEmail,
     cloudContext: CloudContext,
     diskName: DiskName,
     samResource: PersistentDiskSamResourceId,
@@ -581,7 +581,7 @@ class RuntimeV2ServiceInterp[F[_]: Parallel](
       Map(
         "diskName" -> diskName.value,
         "cloudContext" -> cloudContext.asString,
-        "creator" -> userInfo.userEmail.value
+        "creator" -> userEmail.value
       )
 
     // combine default and given labels
@@ -599,10 +599,10 @@ class RuntimeV2ServiceInterp[F[_]: Parallel](
       cloudContext,
       ZoneName("unset"),
       diskName,
-      userInfo.userEmail,
+      userEmail,
       samResource,
       DiskStatus.Creating,
-      AuditInfo(userInfo.userEmail, now, None, now),
+      AuditInfo(userEmail, now, None, now),
       req.azureDiskConfig.size.getOrElse(config.defaultDiskSizeGb),
       req.azureDiskConfig.diskType.getOrElse(config.defaultDiskType),
       config.defaultBlockSizeBytes,
@@ -731,7 +731,7 @@ class RuntimeV2ServiceInterp[F[_]: Parallel](
     workspaceId: WorkspaceId,
     runtimeName: RuntimeName,
     cloudContext: CloudContext,
-    userInfo: UserInfo,
+    userEmail: WorkbenchEmail,
     request: CreateAzureRuntimeRequest,
     samResourceId: RuntimeSamResourceId,
     runtimeImages: Set[RuntimeImage],
@@ -743,9 +743,9 @@ class RuntimeV2ServiceInterp[F[_]: Parallel](
       runtimeName,
       None,
       cloudContext,
-      userInfo.userEmail,
+      userEmail,
       // TODO: use an azure service account
-      Some(userInfo.userEmail),
+      Some(userEmail),
       None,
       None,
       // TODO: Will need to be updated when we support RStudio on Azure or JupyterLab on GCP V2 endpoint
@@ -761,9 +761,9 @@ class RuntimeV2ServiceInterp[F[_]: Parallel](
       runtimeName = runtimeName,
       cloudContext = cloudContext,
       // TODO: use an azure service account
-      serviceAccount = userInfo.userEmail,
+      serviceAccount = userEmail,
       asyncRuntimeFields = None,
-      auditInfo = AuditInfo(userInfo.userEmail, now, None, now),
+      auditInfo = AuditInfo(userEmail, now, None, now),
       kernelFoundBusyDate = None,
       proxyUrl = Runtime.getProxyUrl(config.proxyUrlBase, cloudContext, runtimeName, runtimeImages, None, allLabels),
       status = RuntimeStatus.PreCreating,

diff --git a/...scala/org/broadinstitute/dsde/workbench/leonardo/monitor/LeoPubsubMessageSubscriber.scala b/...scala/org/broadinstitute/dsde/workbench/leonardo/monitor/LeoPubsubMessageSubscriber.scala
@@ -867,13 +867,12 @@ class LeoPubsubMessageSubscriber[F[_]](
             operationAttempt <- F.blocking(v.get()).attempt
             _ <- operationAttempt match {
               case Left(error: java.util.concurrent.ExecutionException) if error.getMessage.contains("Not Found") =>
-                logger.info(ctx.loggingCtx)("disk is already deleted") >> persistentDiskQuery
-                  .delete(diskId, ctx.now)
-                  .transaction[F]
-                  .void >> samService
-                  .getPetServiceAccountToken(disk.auditInfo.creator, googleProject)
-                  .flatMap(petToken => samService.deleteResource(petToken, disk.samResource))
-
+                for {
+                  _ <- logger.info(ctx.loggingCtx)("disk is already deleted")
+                  _ <- persistentDiskQuery.delete(diskId, ctx.now).transaction[F].void
+                  petToken <- samService.getPetServiceAccountToken(disk.auditInfo.creator, googleProject)
+                  _ <- samService.deleteResource(petToken, disk.samResource)
+                } yield ()
               case Left(error) =>
                 F.raiseError(
                   new RuntimeException(s"fail to delete disk ${googleProject}/${disk.name}", error)
@@ -883,9 +882,9 @@ class LeoPubsubMessageSubscriber[F[_]](
                   _ <- F.raiseUnless(isSuccess(op.getHttpErrorStatusCode))(
                     new RuntimeException(s"fail to delete disk ${googleProject}/${disk.name} due to ${op}")
                   )
-                  _ <- persistentDiskQuery.delete(diskId, ctx.now).transaction[F].void >> samService
-                    .getPetServiceAccountToken(disk.auditInfo.creator, googleProject)
-                    .flatMap(petToken => samService.deleteResource(petToken, disk.samResource))
+                  _ <- persistentDiskQuery.delete(diskId, ctx.now).transaction[F].void
+                  petToken <- samService.getPetServiceAccountToken(disk.auditInfo.creator, googleProject)
+                  _ <- samService.deleteResource(petToken, disk.samResource)
                 } yield ()
             }
           } yield ()
@@ -1337,10 +1336,11 @@ class LeoPubsubMessageSubscriber[F[_]](
               // If the message is resubmitted, and this step has already been run, we don't want to re-notify the app creator and update the deleted timestamp
               case AppStatus.Deleted => F.unit
               case _ =>
-                appQuery.markAsDeleted(msg.appId, ctx.now).transaction.void >>
-                  samService
-                    .getPetServiceAccountToken(dbApp.app.auditInfo.creator, msg.project)
-                    .flatMap(petToken => samService.deleteResource(petToken, dbApp.app.samResourceId))
+                for {
+                  _ <- appQuery.markAsDeleted(msg.appId, ctx.now).transaction.void
+                  petToken <- samService.getPetServiceAccountToken(dbApp.app.auditInfo.creator, msg.project)
+                  _ <- samService.deleteResource(petToken, dbApp.app.samResourceId)
+                } yield ()
             }
           else F.unit
       } yield ()

diff --git a/http/src/test/scala/org/broadinstitute/dsde/workbench/leonardo/CommonTestData.scala b/http/src/test/scala/org/broadinstitute/dsde/workbench/leonardo/CommonTestData.scala
@@ -93,16 +93,20 @@ object CommonTestData {
   val userEmail2 = WorkbenchEmail("[email protected]")
   val userEmail3 = WorkbenchEmail("[email protected]")
   val userEmail4 = WorkbenchEmail("[email protected]")
-  val userInfo = UserInfo(OAuth2BearerToken("accessToken"), WorkbenchUserId("user1"), userEmail, 0)
-  val userInfo2 = UserInfo(OAuth2BearerToken("accessToken"), WorkbenchUserId("user2"), userEmail2, 0)
-  val userInfo3 = UserInfo(OAuth2BearerToken("accessToken"), WorkbenchUserId("user3"), userEmail3, 0)
-  val userInfo4 = UserInfo(OAuth2BearerToken("accessToken"), WorkbenchUserId("user4"), userEmail4, 0)
+  val token = OAuth2BearerToken("accessToken1")
+  val token2 = OAuth2BearerToken("accessToken2")
+  val token3 = OAuth2BearerToken("accessToken3")
+  val token4 = OAuth2BearerToken("accessToken4")
+  val userInfo = UserInfo(token, WorkbenchUserId("user1"), userEmail, 0)
+  val userInfo2 = UserInfo(token2, WorkbenchUserId("user2"), userEmail2, 0)
+  val userInfo3 = UserInfo(token3, WorkbenchUserId("user3"), userEmail3, 0)
+  val userInfo4 = UserInfo(token4, WorkbenchUserId("user4"), userEmail4, 0)
   val serviceAccountEmail = WorkbenchEmail("[email protected]")
   val managedIdentityEmail = WorkbenchEmail("subscriptions/foo/bar/userAssignedManagedIdentity")
   val proxyGroupEmail = WorkbenchEmail("[email protected]")
   val unauthorizedEmail = WorkbenchEmail("[email protected]")
   val unauthorizedUserInfo =
-    UserInfo(OAuth2BearerToken("accessToken"), WorkbenchUserId("somecreep"), unauthorizedEmail, 0)
+    UserInfo(OAuth2BearerToken("unauthorized"), WorkbenchUserId("somecreep"), unauthorizedEmail, 0)
   val credentials =
     GoogleCredentials.create(new AccessToken("accessToken", new Date(1000000000000000L))) // don't refresh this token
   val jupyterExtensionBucket = GcsBucketName("bucket-name")