Skip to content

Commit

Permalink
start putting it all together
Browse files Browse the repository at this point in the history
  • Loading branch information
@calypsomatic
calypsomatic committed Nov 7, 2024
1 parent 775616e commit 414207c
Show file tree
Hide file tree
Showing 2 changed files with 235 additions and 29 deletions.
61 changes: 33 additions & 28 deletions core/src/main/scala/org/broadinstitute/dsde/rawls/spendreporting/SpendReportingService.scala
View file
Original file line number Diff line number Diff line change
Expand Up @@ -467,33 +467,34 @@ class SpendReportingService(
Future.failed(RawlsExceptionWithErrorReport(ErrorReport(StatusCodes.InternalServerError, ex)))
}

// def getSpendForAllWorkspaces(
// project: RawlsBillingProjectName,
// start: DateTime,
// end: DateTime,
// aggregations: Set[SpendReportingAggregationKeyWithSub]
// ): Future[SpendReportingResults] = {
// validateReportParameters(start, end)
// requireProjectAction(project, SamBillingProjectActions.readSpendReport) {
// for {
// workspaces <- getOwnerWorkspaces()
// billing <- getBillingForWorkspaces(workspaces)
// query = getAllUserWorkspaceQuery(billing)
// queryJob = setUpAllUserWorkspaceQuery(query, start, end)
//
// job: Job <- bigQueryService.use(_.runJob(queryJob)).unsafeToFuture().map(_.waitFor())
// _ = logSpendQueryStats(job.getStatistics[JobStatistics.QueryStatistics])
// result = job.getQueryResults()
// } yield result.getValues.asScala.toList match {
// case Nil =>
// throw RawlsExceptionWithErrorReport(
// StatusCodes.NotFound,
// s"no spend data found for billing project ${project.value} between dates ${toISODateString(start)} and ${toISODateString(end)}"
// )
// case rows => extractSpendReportingResults(rows, start, end, projectNames, aggregations)
// }
// }
// }
def getSpendForAllWorkspaces(
start: DateTime,
end: DateTime
): Future[SpendReportingResults] = {
validateReportParameters(start, end)
for {
workspaces <- getOwnerWorkspaces()
projectNames = workspaces
.map(wsResp =>
wsResp.workspace.googleProject -> WorkspaceName(wsResp.workspace.namespace, wsResp.workspace.name)
)
.toMap // Map[GoogleProjectId, WorkspaceName]
billing <- getBillingSpendExportsForWorkspaces(workspaces)
query = getAllUserWorkspaceQuery(billing)
queryJob = setUpAllUserWorkspaceQuery(query, start, end)

job: Job <- bigQueryService.use(_.runJob(queryJob)).unsafeToFuture().map(_.waitFor())
_ = logSpendQueryStats(job.getStatistics[JobStatistics.QueryStatistics])
result = job.getQueryResults()
} yield result.getValues.asScala.toList match {
case Nil =>
throw RawlsExceptionWithErrorReport(
StatusCodes.NotFound,
s"no spend data found between dates ${toISODateString(start)} and ${toISODateString(end)}"
) // TODO update this
case rows => extractSpendReportingResults(rows, start, end, projectNames, Set.empty)
}
}

def getOwnerWorkspaces(): Future[Seq[WorkspaceListResponse]] =
workspaceServiceConstructor(ctx).listWorkspaces(WorkspaceFieldSpecs(), -1) map {
Expand All @@ -503,11 +504,15 @@ class SpendReportingService(
case _ => throw new IllegalArgumentException("Expected a JsArray")
}

def getBillingForWorkspaces(
def getBillingSpendExportsForWorkspaces(
workspaces: Seq[WorkspaceListResponse]
): Future[Map[BillingProjectSpendExport, Seq[GoogleProjectId]]] = {
val groupedWorkspaces = workspaces.groupBy(_.workspace.namespace)
val billingProjects = groupedWorkspaces.keys.map(RawlsBillingProjectName).toList
// billingProjects.map(project =>
// requireProjectAction(project, SamBillingProjectActions.readSpendReport)
// ) // TODO a better way to handle this?

getSpendExportConfigurations(billingProjects).map { exportConfigs =>
exportConfigs.map { config =>
config -> groupedWorkspaces(config.billingProjectName.value).map(ws => ws.workspace.googleProject)
Expand Down
203 changes: 202 additions & 1 deletion ...c/test/scala/org/broadinstitute/dsde/rawls/spendreporting/SpendReportingServiceSpec.scala
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1286,7 +1286,7 @@ class SpendReportingServiceSpec extends AnyFlatSpecLike with Matchers with Mocki
)

val result = Await.result(
service.getBillingForWorkspaces(
service.getBillingSpendExportsForWorkspaces(
Seq(workspace2Response, workspace3Response, workspace1Response, workspace4Response)
),
Duration.Inf
Expand Down Expand Up @@ -1432,4 +1432,205 @@ class SpendReportingServiceSpec extends AnyFlatSpecLike with Matchers with Mocki

}

"getSpendForAllWorkspaces" should "get the spend report from multiple billing projects" in {
val from = DateTime.now().minusMonths(2)
val to = from.plusMonths(1)

val price1 = BigDecimal("10.22")
val price2 = BigDecimal("50.74")
val currency = "USD"

val samDAO = mock[SamDAO](RETURNS_SMART_NULLS)
val billingRepository = mock[BillingRepository](RETURNS_SMART_NULLS)
val bpmDAO = mock[BillingProfileManagerDAO](RETURNS_SMART_NULLS)

// Billing projects
val billingProfileId1 = UUID.randomUUID()
val projectName1 = RawlsBillingProjectName("billingProject1")
val billingAccount1 = RawlsBillingAccountName("billingAcct1")
val billingProject1 = RawlsBillingProject(
projectName1,
CreationStatuses.Ready,
Option(billingAccount1),
None,
billingProfileId = Option.apply(billingProfileId1.toString)
)
val billingProfileId2 = UUID.randomUUID()
val projectName2 = RawlsBillingProjectName("billingProject2")
val billingAccount2 = RawlsBillingAccountName("billingAcct2")
val billingProject2 = RawlsBillingProject(
projectName2,
CreationStatuses.Ready,
Option(billingAccount2),
None,
billingProfileId = Option.apply(billingProfileId2.toString)
)
val billingProfileId3 = UUID.randomUUID()
val projectName3 = RawlsBillingProjectName("billingProject3")
val billingAccount3 = RawlsBillingAccountName("billingAcct3")
val billingProject3 = RawlsBillingProject(
projectName3,
CreationStatuses.Ready,
Option(billingAccount3),
None,
billingProfileId = Option.apply(billingProfileId3.toString)
)

// Billing project spend exports
val billingProject1SpendExport =
BillingProjectSpendExport(RawlsBillingProjectName("billingProject1"),
RawlsBillingAccountName("billingAccount1"),
Some("billing1_bq_project.billing1_dataset.billing1_table")
)

val billingProject3SpendExport =
BillingProjectSpendExport(RawlsBillingProjectName("billingProject3"),
RawlsBillingAccountName("billingAccount3"),
None
)

// Workspaces
val workspace1Billing1 =
TestData.workspace("workspace1Billing1",
GoogleProjectId("workspace1ProjectId"),
WorkspaceVersions.V1,
"billingProject1"
)
val workspace2Billing1 =
TestData.workspace("workspace2Billing1",
GoogleProjectId("workspace2ProjectId"),
WorkspaceVersions.V2,
"billingProject1"
)
val workspace1Billing2 =
TestData.workspace("workspace1Billing2",
GoogleProjectId("workspace3ProjectId"),
WorkspaceVersions.V2,
"billingProject2"
)
val workspace1Billing3 =
TestData.workspace("workspace1Billing3",
GoogleProjectId("workspace4ProjectId"),
WorkspaceVersions.V2,
"billingProject3"
)

// Only workspaces 1 and 4 are owned
val workspace1Response = WorkspaceListResponse(
WorkspaceAccessLevels.Owner,
Some(true),
Some(true),
WorkspaceDetails.fromWorkspaceAndOptions(workspace1Billing1,
Option(Set.empty),
true,
Some(WorkspaceCloudPlatform.Gcp)
),
Option.empty,
false,
Some(List.empty)
)
val workspace2Response = WorkspaceListResponse(
WorkspaceAccessLevels.Read,
Some(true),
Some(true),
WorkspaceDetails.fromWorkspaceAndOptions(workspace2Billing1,
Option(Set.empty),
true,
Some(WorkspaceCloudPlatform.Gcp)
),
Option.empty,
false,
Some(List.empty)
)
val workspace3Response = WorkspaceListResponse(
WorkspaceAccessLevels.Write,
Some(true),
Some(true),
WorkspaceDetails.fromWorkspaceAndOptions(workspace1Billing2,
Option(Set.empty),
true,
Some(WorkspaceCloudPlatform.Gcp)
),
Option.empty,
false,
Some(List.empty)
)
val workspace4Response = WorkspaceListResponse(
WorkspaceAccessLevels.Owner,
Some(true),
Some(true),
WorkspaceDetails.fromWorkspaceAndOptions(workspace1Billing3,
Option(Set.empty),
true,
Some(WorkspaceCloudPlatform.Gcp)
),
Option.empty,
false,
Some(List.empty)
)

val dataSource = mock[SlickDataSource]
val mockWorkspaceService = mock[WorkspaceService](RETURNS_SMART_NULLS)

when(mockWorkspaceService.listWorkspaces(any(), any()))
.thenReturn(
Future.successful(Seq(workspace1Response, workspace2Response, workspace3Response, workspace4Response).toJson)
)
val mockWorkspaceServiceConstructor: RawlsRequestContext => WorkspaceService = { _ =>
mockWorkspaceService
}

when(billingRepository.getBillingProject(mockitoEq(projectName1)))
.thenReturn(Future.successful(Option.apply(billingProject1)))
when(billingRepository.getBillingProject(mockitoEq(projectName2)))
.thenReturn(Future.successful(Option.apply(billingProject2)))
when(billingRepository.getBillingProject(mockitoEq(projectName3)))
.thenReturn(Future.successful(Option.apply(billingProject3)))

when(samDAO.userHasAction(any(), any(), any(), any())).thenReturn(Future.successful(true))

val bigQueryService = mockBigQuery(List[Map[String, String]]())

val service = spy(
new SpendReportingService(
testContext,
mock[SlickDataSource],
bigQueryService,
billingRepository,
bpmDAO,
samDAO,
spendReportingServiceConfig,
mockWorkspaceServiceConstructor
)
)
doReturn(Future.successful(Seq(billingProject1SpendExport, billingProject3SpendExport)))
.when(service)
.getSpendExportConfigurations(
any()
)

val spendReport =
TestData.BpmSpendReport.spendData(from, to, currency, Map("Compute" -> price1, "Storage" -> price2))

val billingProfileIdCapture: ArgumentCaptor[UUID] = ArgumentCaptor.forClass(classOf[UUID])
val startDateCapture: ArgumentCaptor[Date] = ArgumentCaptor.forClass(classOf[Date])
val endDateCapture: ArgumentCaptor[Date] = ArgumentCaptor.forClass(classOf[Date])

val result = Await.result(
service.getSpendForAllWorkspaces(from, to),
Duration.Inf
)

result.spendSummary.credits shouldBe "0"
result.spendSummary.cost shouldBe Seq(price1, price2).sum.toString()
result.spendSummary.currency shouldBe "USD"
result.spendSummary.startTime.get.toString(ISODateTimeFormat.date()) shouldBe from.toString(
ISODateTimeFormat.date()
)
result.spendSummary.endTime.get.toString(ISODateTimeFormat.date()) shouldBe to.toString(ISODateTimeFormat.date())

startDateCapture.getValue shouldBe from.toDate
endDateCapture.getValue shouldBe to.toDate
}

}

0 comments on commit 414207c

Please sign in to comment.