diff --git a/common/openapi.yml b/common/openapi.yml index 49ac4023..6ea1a8c5 100644 --- a/common/openapi.yml +++ b/common/openapi.yml @@ -4,7 +4,7 @@ info: version: 1.0.0 paths: # Admin APIs - /api/admin/v1/pipeline/{pipelineName}: + /api/admin/v1/pipelines/{pipelineName}: parameters: - $ref: '#/components/parameters/PipelineName' get: @@ -48,6 +48,51 @@ paths: 500: $ref: '#/components/responses/ServerError' + /api/admin/v1/quotas/{pipelineName}/{userId}: + parameters: + - $ref: '#/components/parameters/PipelineName' + - $ref: '#/components/parameters/UserId' + get: + summary: Get quota for a given pipeline and user. + tags: [ admin ] + operationId: getQuotaForPipelineAndUser + responses: + 200: + description: Success + content: + application/json: + schema: + $ref: '#/components/schemas/AdminQuota' + 400: + $ref: '#/components/responses/BadRequest' + 403: + $ref: '#/components/responses/PermissionDenied' + 500: + $ref: '#/components/responses/ServerError' + patch: + summary: Update quota limit for a given pipeline and user. + tags: [ admin ] + operationId: updateQuotaLimitForPipelineAndUser + requestBody: + required: true + content: + application/json: + schema: + $ref: '#/components/schemas/UpdateQuotaLimitRequestBody' + responses: + 200: + description: Success + content: + application/json: + schema: + $ref: '#/components/schemas/AdminQuota' + 400: + $ref: '#/components/responses/BadRequest' + 403: + $ref: '#/components/responses/PermissionDenied' + 500: + $ref: '#/components/responses/ServerError' + # Pipeline related APIs /api/pipelines/v1: get: @@ -355,6 +400,14 @@ components: schema: type: string + UserId: + name: userId + in: path + description: A string identifier to used to identify a Terra user + required: true + schema: + type: string + WorkspaceId: name: workspaceId in: path @@ -528,6 +581,21 @@ components: workspaceGoogleProject: $ref: "#/components/schemas/PipelineWorkspaceGoogleProject" + AdminQuota: + description: | + Object containing the use id, pipeline identifier, quota limit, and quota usage of a Pipeline for a user. + type: object + required: [ userId, pipelineName, quotaLimit, quotaConsumed ] + properties: + userId: + $ref: "#/components/schemas/UserId" + pipelineName: + $ref: "#/components/schemas/PipelineName" + quotaLimit: + $ref: "#/components/schemas/QuotaLimit" + quotaConsumed: + $ref: "#/components/schemas/QuotaConsumed" + AsyncPipelineRunResponse: description: Result of an asynchronous pipeline run request. type: object @@ -859,6 +927,15 @@ components: wdlMethodVersion: $ref: "#/components/schemas/PipelineWdlMethodVersion" + UpdateQuotaLimitRequestBody: + description: | + json object containing the admin provided information used to update a user's quota limit + type: object + required: [ quotaLimit ] + properties: + quotaLimit: + $ref: "#/components/schemas/QuotaLimit" + UserId: description: | The identifier string for the user who submitted a job request. diff --git a/service/src/main/java/bio/terra/pipelines/app/controller/AdminApiController.java b/service/src/main/java/bio/terra/pipelines/app/controller/AdminApiController.java index dcc27211..21fc836a 100644 --- a/service/src/main/java/bio/terra/pipelines/app/controller/AdminApiController.java +++ b/service/src/main/java/bio/terra/pipelines/app/controller/AdminApiController.java @@ -1,16 +1,20 @@ package bio.terra.pipelines.app.controller; +import bio.terra.common.exception.BadRequestException; import bio.terra.common.iam.SamUser; import bio.terra.common.iam.SamUserFactory; import bio.terra.pipelines.app.configuration.external.SamConfiguration; import bio.terra.pipelines.common.utils.PipelinesEnum; import bio.terra.pipelines.db.entities.Pipeline; +import bio.terra.pipelines.db.entities.UserQuota; import bio.terra.pipelines.dependencies.sam.SamService; import bio.terra.pipelines.generated.api.AdminApi; import bio.terra.pipelines.generated.model.*; import bio.terra.pipelines.service.PipelinesService; +import bio.terra.pipelines.service.QuotasService; import io.swagger.annotations.Api; import jakarta.servlet.http.HttpServletRequest; +import java.util.Optional; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import org.springframework.beans.factory.annotation.Autowired; @@ -27,6 +31,7 @@ public class AdminApiController implements AdminApi { private final HttpServletRequest request; private final PipelinesService pipelinesService; private final SamService samService; + private final QuotasService quotasService; @Autowired public AdminApiController( @@ -34,12 +39,14 @@ public AdminApiController( SamUserFactory samUserFactory, HttpServletRequest request, PipelinesService pipelinesService, - SamService samService) { + SamService samService, + QuotasService quotasService) { this.samConfiguration = samConfiguration; this.samUserFactory = samUserFactory; this.request = request; this.pipelinesService = pipelinesService; this.samService = samService; + this.quotasService = quotasService; } private static final Logger logger = LoggerFactory.getLogger(AdminApiController.class); @@ -58,6 +65,28 @@ public ResponseEntity getPipeline(String pipelineName) { return new ResponseEntity<>(pipelineToApiAdminPipeline(updatedPipeline), HttpStatus.OK); } + @Override + public ResponseEntity getQuotaForPipelineAndUser( + String pipelineName, String userId) { + final SamUser authedUser = getAuthenticatedInfo(); + samService.checkAdminAuthz(authedUser); + PipelinesEnum validatedPipelineName = + PipelineApiUtils.validatePipelineName(pipelineName, logger); + + // Check if a row exists for this user + pipeline. Throw an error if it doesn't + // A row should exist if a user has run into a quota issue before. + Optional userQuota = + quotasService.getQuotaForUserAndPipeline(userId, validatedPipelineName); + if (userQuota.isEmpty()) { + throw new BadRequestException( + String.format( + "User quota not found for user %s and pipeline %s", + userId, validatedPipelineName.getValue())); + } + + return new ResponseEntity<>(userQuotaToApiAdminQuota(userQuota.get()), HttpStatus.OK); + } + @Override public ResponseEntity updatePipeline( String pipelineName, ApiUpdatePipelineRequestBody body) { @@ -74,6 +103,30 @@ public ResponseEntity updatePipeline( return new ResponseEntity<>(pipelineToApiAdminPipeline(updatedPipeline), HttpStatus.OK); } + @Override + public ResponseEntity updateQuotaLimitForPipelineAndUser( + String pipelineName, String userId, ApiUpdateQuotaLimitRequestBody body) { + final SamUser authedUser = getAuthenticatedInfo(); + samService.checkAdminAuthz(authedUser); + PipelinesEnum validatedPipelineName = + PipelineApiUtils.validatePipelineName(pipelineName, logger); + + // Check if a row exists for this user + pipeline. Throw an error if it doesn't + // A row should exist if a user has run into a quota issue before. + Optional userQuota = + quotasService.getQuotaForUserAndPipeline(userId, validatedPipelineName); + if (userQuota.isEmpty()) { + throw new BadRequestException( + String.format( + "User quota not found for user %s and pipeline %s", + userId, validatedPipelineName.getValue())); + } + int newQuotaLimit = body.getQuotaLimit(); + UserQuota updatedUserQuota = + quotasService.adminUpdateQuotaLimit(userQuota.get(), newQuotaLimit); + return new ResponseEntity<>(userQuotaToApiAdminQuota(updatedUserQuota), HttpStatus.OK); + } + public ApiAdminPipeline pipelineToApiAdminPipeline(Pipeline pipeline) { return new ApiAdminPipeline() .pipelineName(pipeline.getName().getValue()) @@ -85,4 +138,12 @@ public ApiAdminPipeline pipelineToApiAdminPipeline(Pipeline pipeline) { .workspaceGoogleProject(pipeline.getWorkspaceGoogleProject()) .wdlMethodVersion(pipeline.getWdlMethodVersion()); } + + public ApiAdminQuota userQuotaToApiAdminQuota(UserQuota userQuota) { + return new ApiAdminQuota() + .userId(userQuota.getUserId()) + .pipelineName(userQuota.getPipelineName().getValue()) + .quotaLimit(userQuota.getQuota()) + .quotaConsumed(userQuota.getQuotaConsumed()); + } } diff --git a/service/src/main/java/bio/terra/pipelines/app/controller/QuotasController.java b/service/src/main/java/bio/terra/pipelines/app/controller/QuotasController.java index f4f61bbd..adbc39dc 100644 --- a/service/src/main/java/bio/terra/pipelines/app/controller/QuotasController.java +++ b/service/src/main/java/bio/terra/pipelines/app/controller/QuotasController.java @@ -49,7 +49,8 @@ public ResponseEntity getQuotaForPipeline(String pipelineNa PipelinesEnum validatedPipelineName = PipelineApiUtils.validatePipelineName(pipelineName, logger); UserQuota userQuota = - quotasService.getQuotaForUserAndPipeline(user.getSubjectId(), validatedPipelineName); + quotasService.getOrCreateQuotaForUserAndPipeline( + user.getSubjectId(), validatedPipelineName); return new ResponseEntity<>(quotasToApiQuotaWithDetails(userQuota), HttpStatus.OK); } diff --git a/service/src/main/java/bio/terra/pipelines/service/QuotasService.java b/service/src/main/java/bio/terra/pipelines/service/QuotasService.java index 9ff88322..f6c2acf3 100644 --- a/service/src/main/java/bio/terra/pipelines/service/QuotasService.java +++ b/service/src/main/java/bio/terra/pipelines/service/QuotasService.java @@ -33,10 +33,9 @@ public class QuotasService { * will create a new row in the user quotas table with the default quota for the pipeline. */ @WriteTransaction - public UserQuota getQuotaForUserAndPipeline(String userId, PipelinesEnum pipelineName) { + public UserQuota getOrCreateQuotaForUserAndPipeline(String userId, PipelinesEnum pipelineName) { // try to get the user quota - Optional userQuota = - userQuotasRepository.findByUserIdAndPipelineName(userId, pipelineName); + Optional userQuota = getQuotaForUserAndPipeline(userId, pipelineName); // if the user quota is not found, grab the default pipeline quota and make a new row in user // quotas table if (userQuota.isEmpty()) { @@ -55,6 +54,14 @@ public UserQuota getQuotaForUserAndPipeline(String userId, PipelinesEnum pipelin return userQuota.get(); } + /** + * This method gets the quota for a given user and pipeline. If the user quota does not exist, it + * will return an empty optional. + */ + public Optional getQuotaForUserAndPipeline(String userId, PipelinesEnum pipelineName) { + return userQuotasRepository.findByUserIdAndPipelineName(userId, pipelineName); + } + /** * This method updates the quota consumed for a given user and pipeline. It will return the * updated user quota object. @@ -75,4 +82,23 @@ public UserQuota updateQuotaConsumed(UserQuota userQuota, int newQuotaConsumed) userQuota.setQuotaConsumed(newQuotaConsumed); return userQuotasRepository.save(userQuota); } + + /** + * This method updates the quota limit for a given user quota. This should only be called from the + * Admin Controller + * + * @param userQuota - the user quota to update + * @param newQuotaLimit - the new quota limit + * @return - the updated user quota + */ + public UserQuota adminUpdateQuotaLimit(UserQuota userQuota, int newQuotaLimit) { + if (newQuotaLimit < userQuota.getQuotaConsumed()) { + throw new InternalServerErrorException( + String.format( + "New quota limit: %d, is less than the quota consumed: %d", + newQuotaLimit, userQuota.getQuotaConsumed())); + } + userQuota.setQuota(newQuotaLimit); + return userQuotasRepository.save(userQuota); + } } diff --git a/service/src/main/java/bio/terra/pipelines/stairway/QuotaConsumedValidationStep.java b/service/src/main/java/bio/terra/pipelines/stairway/QuotaConsumedValidationStep.java index bbc4a8ea..da9e01c7 100644 --- a/service/src/main/java/bio/terra/pipelines/stairway/QuotaConsumedValidationStep.java +++ b/service/src/main/java/bio/terra/pipelines/stairway/QuotaConsumedValidationStep.java @@ -43,7 +43,7 @@ public StepResult doStep(FlightContext flightContext) { // check if user quota used plus quota consumed is less than or equal to user quota Integer quotaUsedForThisRun = workingMap.get(ImputationJobMapKeys.QUOTA_CONSUMED, Integer.class); - UserQuota userQuota = quotasService.getQuotaForUserAndPipeline(userId, pipelineName); + UserQuota userQuota = quotasService.getOrCreateQuotaForUserAndPipeline(userId, pipelineName); // user quota has been exceeded, fail the flight int totalQuotaConsumed = userQuota.getQuotaConsumed() + quotaUsedForThisRun; @@ -79,7 +79,7 @@ public StepResult undoStep(FlightContext flightContext) { // update the user quota to be what it was before this run's quota was added int quotaUsedForThisRun = workingMap.get(ImputationJobMapKeys.QUOTA_CONSUMED, Integer.class); - UserQuota userQuota = quotasService.getQuotaForUserAndPipeline(userId, pipelineName); + UserQuota userQuota = quotasService.getOrCreateQuotaForUserAndPipeline(userId, pipelineName); quotasService.updateQuotaConsumed( userQuota, userQuota.getQuotaConsumed() - quotaUsedForThisRun); diff --git a/service/src/test/java/bio/terra/pipelines/controller/AdminApiControllerTest.java b/service/src/test/java/bio/terra/pipelines/controller/AdminApiControllerTest.java index b0a0e10a..eb2ad19c 100644 --- a/service/src/test/java/bio/terra/pipelines/controller/AdminApiControllerTest.java +++ b/service/src/test/java/bio/terra/pipelines/controller/AdminApiControllerTest.java @@ -17,14 +17,19 @@ import bio.terra.pipelines.app.controller.AdminApiController; import bio.terra.pipelines.app.controller.GlobalExceptionHandler; import bio.terra.pipelines.common.utils.PipelinesEnum; +import bio.terra.pipelines.db.entities.UserQuota; import bio.terra.pipelines.dependencies.sam.SamService; import bio.terra.pipelines.generated.model.ApiAdminPipeline; +import bio.terra.pipelines.generated.model.ApiAdminQuota; import bio.terra.pipelines.generated.model.ApiUpdatePipelineRequestBody; +import bio.terra.pipelines.generated.model.ApiUpdateQuotaLimitRequestBody; import bio.terra.pipelines.service.PipelinesService; +import bio.terra.pipelines.service.QuotasService; import bio.terra.pipelines.testutils.MockMvcUtils; import com.fasterxml.jackson.core.JsonProcessingException; import com.fasterxml.jackson.databind.ObjectMapper; import jakarta.servlet.http.HttpServletRequest; +import java.util.Optional; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.springframework.beans.factory.annotation.Autowired; @@ -39,6 +44,7 @@ @WebMvcTest() class AdminApiControllerTest { @MockBean PipelinesService pipelinesServiceMock; + @MockBean QuotasService quotasServiceMock; @MockBean SamUserFactory samUserFactoryMock; @MockBean BearerTokenFactory bearerTokenFactory; @MockBean SamConfiguration samConfiguration; @@ -68,7 +74,8 @@ void updatePipelineWorkspaceOk() throws Exception { .perform( patch( String.format( - "/api/admin/v1/pipeline/%s", PipelinesEnum.ARRAY_IMPUTATION.getValue())) + "/api/admin/v1/pipelines/%s", + PipelinesEnum.ARRAY_IMPUTATION.getValue())) .contentType(MediaType.APPLICATION_JSON) .content( createTestJobPostBody( @@ -99,7 +106,7 @@ void updatePipelineWorkspaceIdRequireWorkspaceName() throws Exception { .perform( patch( String.format( - "/api/admin/v1/pipeline/%s", PipelinesEnum.ARRAY_IMPUTATION.getValue())) + "/api/admin/v1/pipelines/%s", PipelinesEnum.ARRAY_IMPUTATION.getValue())) .contentType(MediaType.APPLICATION_JSON) .content( createTestJobPostBody( @@ -113,7 +120,7 @@ void updatePipelineWorkspaceIdRequireWorkspaceProject() throws Exception { .perform( patch( String.format( - "/api/admin/v1/pipeline/%s", PipelinesEnum.ARRAY_IMPUTATION.getValue())) + "/api/admin/v1/pipelines/%s", PipelinesEnum.ARRAY_IMPUTATION.getValue())) .contentType(MediaType.APPLICATION_JSON) .content(createTestJobPostBody(null, TEST_WORKSPACE_NAME, TEST_WDL_METHOD_VERSION))) .andExpect(status().isBadRequest()); @@ -125,7 +132,7 @@ void updatePipelineWorkspaceIdRequireWdlMethodVersion() throws Exception { .perform( patch( String.format( - "/api/admin/v1/pipeline/%s", PipelinesEnum.ARRAY_IMPUTATION.getValue())) + "/api/admin/v1/pipelines/%s", PipelinesEnum.ARRAY_IMPUTATION.getValue())) .contentType(MediaType.APPLICATION_JSON) .content( createTestJobPostBody( @@ -141,7 +148,7 @@ void updatePipelineWorkspaceIdNotAdminUser() throws Exception { .perform( patch( String.format( - "/api/admin/v1/pipeline/%s", PipelinesEnum.ARRAY_IMPUTATION.getValue())) + "/api/admin/v1/pipelines/%s", PipelinesEnum.ARRAY_IMPUTATION.getValue())) .contentType(MediaType.APPLICATION_JSON) .content( createTestJobPostBody( @@ -160,7 +167,7 @@ void getAdminPipelineOk() throws Exception { .perform( get( String.format( - "/api/admin/v1/pipeline/%s", PipelinesEnum.ARRAY_IMPUTATION.getValue()))) + "/api/admin/v1/pipelines/%s", PipelinesEnum.ARRAY_IMPUTATION.getValue()))) .andExpect(status().isOk()) .andExpect(content().contentType(MediaType.APPLICATION_JSON)) .andReturn(); @@ -182,7 +189,140 @@ void getAdminPipelineNotAdminUser() throws Exception { .perform( get( String.format( - "/api/admin/v1/pipeline/%s", PipelinesEnum.ARRAY_IMPUTATION.getValue()))) + "/api/admin/v1/pipelines/%s", PipelinesEnum.ARRAY_IMPUTATION.getValue()))) + .andExpect(status().isForbidden()); + } + + @Test + void getUserQuotaOk() throws Exception { + when(quotasServiceMock.getQuotaForUserAndPipeline( + TEST_SAM_USER.getSubjectId(), PipelinesEnum.ARRAY_IMPUTATION)) + .thenReturn(Optional.of(TEST_USER_QUOTA_1)); + MvcResult result = + mockMvc + .perform( + get( + String.format( + "/api/admin/v1/quotas/%s/%s", + PipelinesEnum.ARRAY_IMPUTATION.getValue(), TEST_SAM_USER.getSubjectId()))) + .andExpect(status().isOk()) + .andExpect(content().contentType(MediaType.APPLICATION_JSON)) + .andReturn(); + + ApiAdminQuota response = + new ObjectMapper() + .readValue(result.getResponse().getContentAsString(), ApiAdminQuota.class); + + // this is all mocked data so really not worth checking values, really just testing that it's a + // 200 status with a properly formatted response + assertEquals(PipelinesEnum.ARRAY_IMPUTATION.getValue(), response.getPipelineName()); + assertEquals(TEST_USER_QUOTA_1.getUserId(), response.getUserId()); + assertEquals(TEST_USER_QUOTA_1.getQuota(), response.getQuotaLimit()); + assertEquals(TEST_USER_QUOTA_1.getQuotaConsumed(), response.getQuotaConsumed()); + } + + @Test + void getAdminQuotaNotAdminUser() throws Exception { + doThrow(new ForbiddenException("error string")).when(samServiceMock).checkAdminAuthz(testUser); + + mockMvc + .perform( + get( + String.format( + "/api/admin/v1/quotas/%s/%s", + PipelinesEnum.ARRAY_IMPUTATION.getValue(), TEST_SAM_USER.getSubjectId()))) + .andExpect(status().isForbidden()); + } + + @Test + void getAdminQuotaUserQuotaDoesntExist() throws Exception { + when(quotasServiceMock.getQuotaForUserAndPipeline( + TEST_SAM_USER.getSubjectId(), PipelinesEnum.ARRAY_IMPUTATION)) + .thenReturn(Optional.empty()); + mockMvc + .perform( + get( + String.format( + "/api/admin/v1/quotas/%s/%s", + PipelinesEnum.ARRAY_IMPUTATION.getValue(), TEST_SAM_USER.getSubjectId()))) + .andExpect(status().isBadRequest()); + } + + @Test + void updateAdminQuotaOk() throws Exception { + UserQuota updatedUserQuota = + new UserQuota(PipelinesEnum.ARRAY_IMPUTATION, TEST_SAM_USER.getSubjectId(), 800, 0); + when(quotasServiceMock.getQuotaForUserAndPipeline( + TEST_SAM_USER.getSubjectId(), PipelinesEnum.ARRAY_IMPUTATION)) + .thenReturn(Optional.of(TEST_USER_QUOTA_1)); + when(quotasServiceMock.adminUpdateQuotaLimit(TEST_USER_QUOTA_1, 800)) + .thenReturn(updatedUserQuota); + MvcResult result = + mockMvc + .perform( + patch( + String.format( + "/api/admin/v1/quotas/%s/%s", + PipelinesEnum.ARRAY_IMPUTATION.getValue(), + TEST_SAM_USER.getSubjectId())) + .contentType(MediaType.APPLICATION_JSON) + .content(createTestJobPostBody(800))) + .andExpect(status().isOk()) + .andExpect(content().contentType(MediaType.APPLICATION_JSON)) + .andReturn(); + + ApiAdminQuota response = + new ObjectMapper() + .readValue(result.getResponse().getContentAsString(), ApiAdminQuota.class); + + // this is all mocked data so really not worth checking values, really just testing that it's a + // 200 status with a properly formatted response + assertEquals(PipelinesEnum.ARRAY_IMPUTATION.getValue(), response.getPipelineName()); + assertEquals(TEST_SAM_USER.getSubjectId(), response.getUserId()); + assertEquals(800, response.getQuotaLimit()); + } + + @Test + void updateAdminQuotaUserQuotaDoesntExist() throws Exception { + when(quotasServiceMock.getQuotaForUserAndPipeline( + TEST_SAM_USER.getSubjectId(), PipelinesEnum.ARRAY_IMPUTATION)) + .thenReturn(Optional.empty()); + mockMvc + .perform( + patch( + String.format( + "/api/admin/v1/quotas/%s/%s", + PipelinesEnum.ARRAY_IMPUTATION.getValue(), TEST_SAM_USER.getSubjectId())) + .contentType(MediaType.APPLICATION_JSON) + .content(createTestJobPostBody(800))) + .andExpect(status().isBadRequest()); + } + + @Test + void updateAdminQuotaRequireQuotaLimit() throws Exception { + mockMvc + .perform( + patch( + String.format( + "/api/admin/v1/quotas/%s/%s", + PipelinesEnum.ARRAY_IMPUTATION.getValue(), TEST_SAM_USER.getSubjectId())) + .contentType(MediaType.APPLICATION_JSON) + .content("{}")) + .andExpect(status().isBadRequest()); + } + + @Test + void updateAdminQuotaIdNotAdminUser() throws Exception { + doThrow(new ForbiddenException("error string")).when(samServiceMock).checkAdminAuthz(testUser); + + mockMvc + .perform( + patch( + String.format( + "/api/admin/v1/quotas/%s/%s", + PipelinesEnum.ARRAY_IMPUTATION.getValue(), TEST_SAM_USER.getSubjectId())) + .contentType(MediaType.APPLICATION_JSON) + .content(createTestJobPostBody(500))) .andExpect(status().isForbidden()); } @@ -196,4 +336,10 @@ private String createTestJobPostBody( .wdlMethodVersion(wdlMethodVersion); return MockMvcUtils.convertToJsonString(apiUpdatePipelineRequestBody); } + + private String createTestJobPostBody(int quotaLimit) throws JsonProcessingException { + ApiUpdateQuotaLimitRequestBody apiUpdateQuotaLimitRequestBody = + new ApiUpdateQuotaLimitRequestBody().quotaLimit(quotaLimit); + return MockMvcUtils.convertToJsonString(apiUpdateQuotaLimitRequestBody); + } } diff --git a/service/src/test/java/bio/terra/pipelines/controller/QuotasControllerTest.java b/service/src/test/java/bio/terra/pipelines/controller/QuotasControllerTest.java index 89409c23..c981338b 100644 --- a/service/src/test/java/bio/terra/pipelines/controller/QuotasControllerTest.java +++ b/service/src/test/java/bio/terra/pipelines/controller/QuotasControllerTest.java @@ -49,7 +49,7 @@ void beforeEach() { when(samConfiguration.baseUri()).thenReturn("baseSamUri"); when(samUserFactoryMock.from(any(HttpServletRequest.class), eq("baseSamUri"))) .thenReturn(testUser); - when(quotasServiceMock.getQuotaForUserAndPipeline( + when(quotasServiceMock.getOrCreateQuotaForUserAndPipeline( testUser.getSubjectId(), PipelinesEnum.ARRAY_IMPUTATION)) .thenReturn(testUserQuota); } diff --git a/service/src/test/java/bio/terra/pipelines/service/QuotasServiceTest.java b/service/src/test/java/bio/terra/pipelines/service/QuotasServiceTest.java index 17a79be2..dcacbee8 100644 --- a/service/src/test/java/bio/terra/pipelines/service/QuotasServiceTest.java +++ b/service/src/test/java/bio/terra/pipelines/service/QuotasServiceTest.java @@ -19,13 +19,30 @@ class QuotasServiceTest extends BaseEmbeddedDbTest { @Test void getQuotaForUserAndPipeline() { + // add row to user_quotas table + createAndSaveUserQuota(TestUtils.TEST_USER_ID_1, PipelinesEnum.ARRAY_IMPUTATION, 30, 100); + + // call service and assert correct UserQuota is returned + assertTrue( + quotasService + .getQuotaForUserAndPipeline(TestUtils.TEST_USER_ID_1, PipelinesEnum.ARRAY_IMPUTATION) + .isPresent()); + // assert no UserQuota is returned for a non-existing user + assertTrue( + quotasService + .getQuotaForUserAndPipeline("randomUser", PipelinesEnum.ARRAY_IMPUTATION) + .isEmpty()); + } + + @Test + void getOrCreateQuotaForUserAndPipeline() { // add row to user_quotas table UserQuota userQuota = createAndSaveUserQuota(TestUtils.TEST_USER_ID_1, PipelinesEnum.ARRAY_IMPUTATION, 30, 100); // call service and assert correct UserQuota is returned UserQuota returnedUserQuota = - quotasService.getQuotaForUserAndPipeline( + quotasService.getOrCreateQuotaForUserAndPipeline( TestUtils.TEST_USER_ID_1, PipelinesEnum.ARRAY_IMPUTATION); assertEquals(userQuota.getUserId(), returnedUserQuota.getUserId()); @@ -50,7 +67,7 @@ void addRowForNonExistentUserQuota() { // call service with same inputs and a new row should exist in user_quotas table UserQuota userQuota = - quotasService.getQuotaForUserAndPipeline( + quotasService.getOrCreateQuotaForUserAndPipeline( TestUtils.TEST_USER_ID_1, PipelinesEnum.ARRAY_IMPUTATION); assertEquals(TestUtils.TEST_USER_ID_1, userQuota.getUserId()); @@ -116,4 +133,33 @@ void updateQuotaConsumedGreaterThanQuota() { InternalServerErrorException.class, () -> quotasService.updateQuotaConsumed(userQuota, newQuotaConsumed)); } + + @Test + void updateQuotaLimit() { + // add row to user_quotas table + UserQuota userQuota = + createAndSaveUserQuota(TestUtils.TEST_USER_ID_1, PipelinesEnum.ARRAY_IMPUTATION, 30, 100); + + // call service to update quota limit + int newQuotaLimit = 150; + UserQuota updatedUserQuota = quotasService.adminUpdateQuotaLimit(userQuota, newQuotaLimit); + + assertEquals(userQuota.getUserId(), updatedUserQuota.getUserId()); + assertEquals(userQuota.getPipelineName(), updatedUserQuota.getPipelineName()); + assertEquals(userQuota.getQuotaConsumed(), updatedUserQuota.getQuotaConsumed()); + assertEquals(newQuotaLimit, updatedUserQuota.getQuota()); + } + + @Test + void updateQuotaLimitLessThanQuotaConsumed() { + // add row to user_quotas table + UserQuota userQuota = + createAndSaveUserQuota(TestUtils.TEST_USER_ID_1, PipelinesEnum.ARRAY_IMPUTATION, 30, 100); + + // call service to update quota limit + int newQuotaLimit = 20; + assertThrows( + InternalServerErrorException.class, + () -> quotasService.adminUpdateQuotaLimit(userQuota, newQuotaLimit)); + } } diff --git a/service/src/test/java/bio/terra/pipelines/stairway/QuotaConsumedValidationStepTest.java b/service/src/test/java/bio/terra/pipelines/stairway/QuotaConsumedValidationStepTest.java index 2b5ec446..c8028c81 100644 --- a/service/src/test/java/bio/terra/pipelines/stairway/QuotaConsumedValidationStepTest.java +++ b/service/src/test/java/bio/terra/pipelines/stairway/QuotaConsumedValidationStepTest.java @@ -46,7 +46,7 @@ void doStepSuccess() { // before running make sure quota consumed for user is 0 UserQuota userQuota = - quotasService.getQuotaForUserAndPipeline( + quotasService.getOrCreateQuotaForUserAndPipeline( TestUtils.TEST_USER_ID_1, PipelinesEnum.ARRAY_IMPUTATION); assertEquals(0, userQuota.getQuotaConsumed()); @@ -60,7 +60,7 @@ void doStepSuccess() { // after running make sure quota for user is 30 userQuota = - quotasService.getQuotaForUserAndPipeline( + quotasService.getOrCreateQuotaForUserAndPipeline( TestUtils.TEST_USER_ID_1, PipelinesEnum.ARRAY_IMPUTATION); assertEquals(30, userQuota.getQuotaConsumed()); } @@ -98,7 +98,7 @@ void undoStepSuccess() { assertEquals(StepStatus.STEP_RESULT_SUCCESS, result.getStepStatus()); // the working map has 30 quota consumed, so the user quota consumed should be 50 - 30 = 20 UserQuota userQuota = - quotasService.getQuotaForUserAndPipeline( + quotasService.getOrCreateQuotaForUserAndPipeline( TestUtils.TEST_USER_ID_1, PipelinesEnum.ARRAY_IMPUTATION); assertEquals(20, userQuota.getQuotaConsumed()); }