temp

Author: sluongng

enterprise/server/test/integration/usage/BUILD

@@ -0,0 +1,19 @@
+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:usage_go_proto",
+        "//server/tables",
+        "@com_github_stretchr_testify//assert",
+        "@com_github_stretchr_testify//require",
+        "@org_golang_google_protobuf//encoding/protojson",
+    ],
+)

enterprise/server/test/integration/usage/usage_test.go

@@ -0,0 +1,123 @@
+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"
+	usagepb "github.com/buildbuddy-io/buildbuddy/proto/usage"
+)
+
+func TestGetUsage_ParentAdminFetchesChildGroupUsage(t *testing.T) {
+	// Start enterprise app with usage tracking enabled.
+	app := buildbuddy_enterprise.Run(t, "--app.usage_tracking_enabled=true", "--app.region=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 a child group (same SAML as parent). Do NOT add the user as a member.
+	child := &tables.Group{Name: "child-org", URLIdentifier: "child-org", SamlIdpMetadataUrl: parent.SamlIdpMetadataUrl}
+	require.NoError(t, app.DB().Create(child).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",
+		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)
+
+	// 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)
+	periodStart := monthStart.Add(6 * time.Hour)
+	seeded := &tables.Usage{
+		GroupID:         child.GroupID,
+		PeriodStartUsec: periodStart.UnixMicro(),
+		Region:          "test",
+		FinalBeforeUsec: periodStart.Add(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.
+	// Example curl:
+	//   curl -X GET \
+	//     -H "Content-Type: application/json" \
+	//     -H "x-buildbuddy-api-key: <PARENT_ADMIN_API_KEY>" \
+	//     --data '{"requestContext":{"groupId":"<CHILD_GROUP_ID>"}}' \
+	//     "http://localhost:8080/rpc/BuildBuddyService/GetUsage"
+	// Example body for users:
+	// {
+	//   "requestContext": {
+	//     "groupId": "<child-group-id>"
+	//   }
+	// }
+	jsonBody := fmt.Sprintf(`
+	{
+		"requestContext": {
+			"group_id": "%s"
+		}
+	}`, child.GroupID)
+	httpReq, err := http.NewRequest(http.MethodPost, app.HTTPURL()+"/rpc/BuildBuddyService/GetUsage", bytes.NewBufferString(jsonBody))
+	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, 200, 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())
+}