diff --git a/enterprise/server/test/integration/usage/BUILD b/enterprise/server/test/integration/usage/BUILD new file mode 100644 index 00000000000..645ea3be00d --- /dev/null +++ b/enterprise/server/test/integration/usage/BUILD @@ -0,0 +1,20 @@ +load("@io_bazel_rules_go//go:def.bzl", "go_test") + +package(default_visibility = ["//enterprise:__subpackages__"]) + +go_test( + name = "usage_test", + srcs = ["usage_test.go"], + deps = [ + "//enterprise/server/testutil/buildbuddy_enterprise", + "//proto:api_key_go_proto", + "//proto:capability_go_proto", + "//proto:context_go_proto", + "//proto:group_go_proto", + "//proto:usage_go_proto", + "//server/tables", + "@com_github_stretchr_testify//assert", + "@com_github_stretchr_testify//require", + "@org_golang_google_protobuf//encoding/protojson", + ], +) diff --git a/enterprise/server/test/integration/usage/usage_test.go b/enterprise/server/test/integration/usage/usage_test.go new file mode 100644 index 00000000000..a86a7013e02 --- /dev/null +++ b/enterprise/server/test/integration/usage/usage_test.go @@ -0,0 +1,131 @@ +package usage_test + +import ( + "bytes" + "fmt" + "io" + "net/http" + "testing" + "time" + + "github.com/buildbuddy-io/buildbuddy/enterprise/server/testutil/buildbuddy_enterprise" + "github.com/buildbuddy-io/buildbuddy/server/tables" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "google.golang.org/protobuf/encoding/protojson" + + akpb "github.com/buildbuddy-io/buildbuddy/proto/api_key" + cappb "github.com/buildbuddy-io/buildbuddy/proto/capability" + ctxpb "github.com/buildbuddy-io/buildbuddy/proto/context" + grpb "github.com/buildbuddy-io/buildbuddy/proto/group" + usagepb "github.com/buildbuddy-io/buildbuddy/proto/usage" +) + +func TestGetUsage_ParentAdminFetchesChildGroupUsage(t *testing.T) { + // Start enterprise app + app := buildbuddy_enterprise.Run( + t, + "--app.usage_tracking_enabled=true", + "--app.region=test", // Required for usage tracking + "--auth.api_key_group_cache_ttl=0", // Don't cache API key -> group mappings for this test + ) + wc := buildbuddy_enterprise.LoginAsDefaultSelfAuthUser(t, app) + + // Treat the logged-in user's selected group as the parent org; mark it as parent and + // set SAML metadata so that authdb recognizes child orgs sharing the same IdP. + parentGroupID := wc.RequestContext.GetGroupId() + parent := &tables.Group{} + require.NoError(t, app.DB().Where("group_id = ?", parentGroupID).Take(parent).Error) + parent.IsParent = true + parent.SamlIdpMetadataUrl = "https://idp.example.test/metadata" + parent.URLIdentifier = "parent-org" + require.NoError(t, app.DB().Save(parent).Error) + + // Create an ORG_ADMIN API key for the parent org using cookie-authenticated Web RPC. + createKeyReq := &akpb.CreateApiKeyRequest{ + RequestContext: &ctxpb.RequestContext{GroupId: parentGroupID}, + Label: "admin-1", + Capability: []cappb.Capability{cappb.Capability_ORG_ADMIN}, + } + createKeyRsp := &akpb.CreateApiKeyResponse{} + require.NoError(t, wc.RPC("CreateApiKey", createKeyReq, createKeyRsp)) + adminAPIKey := createKeyRsp.GetApiKey().GetValue() + require.NotEmpty(t, adminAPIKey) + + // Create a child group (same SAML as parent) using the API. + // Do NOT add the user as a member: authenticate with the parent ORG_ADMIN API key. + createChildGrpReq, err := http.NewRequest( + http.MethodPost, + app.HTTPURL()+"/rpc/BuildBuddyService/CreateGroup", + bytes.NewBufferString(`{"name":"child-org","urlIdentifier":"child-org"}`), + ) + require.NoError(t, err) + createChildGrpReq.Header.Set("Content-Type", "application/json") + createChildGrpReq.Header.Set("x-buildbuddy-api-key", adminAPIKey) + createChildGrpResp, err := http.DefaultClient.Do(createChildGrpReq) + require.NoError(t, err) + defer createChildGrpResp.Body.Close() + createResBytes, err := io.ReadAll(createChildGrpResp.Body) + require.NoError(t, err) + require.Equal(t, http.StatusOK, createChildGrpResp.StatusCode) + childGrpRsp := &grpb.CreateGroupResponse{} + require.NoError(t, protojson.Unmarshal(createResBytes, childGrpRsp)) + childGroupID := childGrpRsp.GetId() + require.NotEmpty(t, childGroupID) + child := &tables.Group{} + require.NoError(t, app.DB().Where("group_id = ?", childGroupID).Take(child).Error) + + // Seed a usage row for the child group. + now := time.Now().UTC() + monthStart := time.Date(now.Year(), now.Month(), 1, 0, 0, 0, 0, time.UTC) + seeded := &tables.Usage{ + GroupID: childGroupID, + PeriodStartUsec: monthStart.UnixMicro(), + Region: "test", + FinalBeforeUsec: monthStart.Add(30 * 24 * time.Hour).UnixMicro(), + UsageCounts: tables.UsageCounts{ + Invocations: 3, + CASCacheHits: 1, + ActionCacheHits: 2, + TotalDownloadSizeBytes: 2048, + LinuxExecutionDurationUsec: 111, + TotalUploadSizeBytes: 4096, + TotalCachedActionExecUsec: 222, + CPUNanos: 333, + }, + UsageLabels: tables.UsageLabels{Origin: "internal", Client: "bazel", Server: "app"}, + } + db := app.DB() + require.NoError(t, db.Create(seeded).Error) + + // Make the HTTP protolet request using the parent admin API key and requesting the child group. + // Use JSON body for documentation clarity. + httpReq, err := http.NewRequest( + http.MethodPost, + app.HTTPURL()+"/rpc/BuildBuddyService/GetUsage", + bytes.NewBufferString(fmt.Sprintf(`{"requestContext":{"groupId":"%s"}}`, childGroupID)), + ) + require.NoError(t, err) + httpReq.Header.Set("Content-Type", "application/json") + httpReq.Header.Set("x-buildbuddy-api-key", adminAPIKey) + httpRes, err := http.DefaultClient.Do(httpReq) + require.NoError(t, err) + defer httpRes.Body.Close() + resBytes, err := io.ReadAll(httpRes.Body) + require.NoError(t, err) + require.Equal(t, http.StatusOK, httpRes.StatusCode) + rsp := &usagepb.GetUsageResponse{} + require.NoError(t, protojson.Unmarshal(resBytes, rsp)) + + require.NotNil(t, rsp.GetUsage()) + assert.EqualValues(t, 3, rsp.GetUsage().GetInvocations()) + assert.EqualValues(t, 2, rsp.GetUsage().GetActionCacheHits()) + assert.EqualValues(t, 1, rsp.GetUsage().GetCasCacheHits()) + assert.EqualValues(t, 2048, rsp.GetUsage().GetTotalDownloadSizeBytes()) + assert.EqualValues(t, 4096, rsp.GetUsage().GetTotalUploadSizeBytes()) + assert.EqualValues(t, 111, rsp.GetUsage().GetLinuxExecutionDurationUsec()) + assert.EqualValues(t, 222, rsp.GetUsage().GetTotalCachedActionExecUsec()) + assert.EqualValues(t, 333, rsp.GetUsage().GetCloudCpuNanos()) + require.Len(t, rsp.GetDailyUsage(), 1) + assert.EqualValues(t, 3, rsp.GetDailyUsage()[0].GetInvocations()) +}