From dce377673776c16db8022436fbdb72b6055d4748 Mon Sep 17 00:00:00 2001
From: Eric Rollins <erollins@broadinstitute.org>
Date: Wed, 20 Nov 2024 12:26:51 -0500
Subject: [PATCH] Update workspace status when credits are extended.

---
 .../initialcredits/InitialCreditsService.java | 37 ++++++++++++-------
 1 file changed, 23 insertions(+), 14 deletions(-)

diff --git a/api/src/main/java/org/pmiops/workbench/initialcredits/InitialCreditsService.java b/api/src/main/java/org/pmiops/workbench/initialcredits/InitialCreditsService.java
index f353cc32893..c4e65c9f3d1 100644
--- a/api/src/main/java/org/pmiops/workbench/initialcredits/InitialCreditsService.java
+++ b/api/src/main/java/org/pmiops/workbench/initialcredits/InitialCreditsService.java
@@ -375,6 +375,9 @@ public DbUser extendInitialCreditsExpiration(DbUser user) {
                 + TimeUnit.DAYS.toMillis(
                     workbenchConfigProvider.get().billing.initialCreditsExtensionPeriodDays)));
     userInitialCreditsExpiration.setExtensionTime(clockNow());
+
+    updateExpirationOfWorkspacesOwnedByUser(user, false);
+
     return userDao.save(user);
   }
 
@@ -419,20 +422,7 @@ private void handleExpiredCredits(
         user.getUsername(),
         userInitialCreditsExpiration.getExpirationTime());
 
-    workspaceDao.findAllByCreator(user).stream()
-        .filter(
-            ws ->
-                BillingUtils.isInitialCredits(
-                    ws.getBillingAccountName(), workbenchConfigProvider.get()))
-        .filter(DbWorkspace::isActive)
-        .filter(ws -> !ws.isInitialCreditsExpired())
-        .forEach(
-            ws -> {
-              ws.setInitialCreditsExpired(true);
-              ws.setBillingStatus(BillingStatus.INACTIVE);
-              workspaceDao.save(ws);
-              deleteAppsAndRuntimesInWorkspace(ws);
-            });
+    updateExpirationOfWorkspacesOwnedByUser(user, true);
 
     userInitialCreditsExpiration.setExpirationCleanupTime(clockNow());
     userDao.save(user);
@@ -674,4 +664,23 @@ private static Map<Long, Double> getDbCostByCreatorCache(
                         v -> Optional.ofNullable(v.getFreeTierCost()).orElse(0.0))));
     return dbCostByCreator;
   }
+
+  private void updateExpirationOfWorkspacesOwnedByUser(DbUser user, boolean isExpired) {
+    workspaceDao.findAllByCreator(user).stream()
+        .filter(
+            ws ->
+                BillingUtils.isInitialCredits(
+                    ws.getBillingAccountName(), workbenchConfigProvider.get()))
+        .filter(DbWorkspace::isActive)
+        .filter(ws -> ws.isInitialCreditsExpired() != isExpired)
+        .forEach(
+            ws -> {
+              ws.setInitialCreditsExpired(isExpired);
+              ws.setBillingStatus(isExpired ? BillingStatus.INACTIVE : BillingStatus.ACTIVE);
+              workspaceDao.save(ws);
+              if (isExpired) {
+                deleteAppsAndRuntimesInWorkspace(ws);
+              }
+            });
+  }
 }