From 1ea835cefff07f329218e6e7c780e73196005c43 Mon Sep 17 00:00:00 2001 From: purplenicole730 Date: Wed, 20 Nov 2024 10:33:01 -0500 Subject: [PATCH 01/15] add regular billing wrappers --- app/billing_client.go | 234 ++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 234 insertions(+) create mode 100644 app/billing_client.go diff --git a/app/billing_client.go b/app/billing_client.go new file mode 100644 index 00000000000..529ff3e1948 --- /dev/null +++ b/app/billing_client.go @@ -0,0 +1,234 @@ +package app + +import ( + "context" + + pb "go.viam.com/api/app/v1" + "go.viam.com/utils/rpc" + "google.golang.org/protobuf/types/known/timestamppb" +) + +type UsageCostType int32 + +const ( + UsageCostTypeUnspecified UsageCostType = iota + UsageCostTypeDataUpload + UsageCostTypeDataEgress + UsageCostTypeRemoteControl + UsageCostTypeStandardCompute + UsageCostTypeCloudStorage + UsageCostTypeBinaryDataCloudStorage + UsageCostTypeOtherCloudStorage + UsageCostTypePerMachine +) + +func usageCostTypeFromProto(costType pb.UsageCostType) UsageCostType { + switch costType { + case pb.UsageCostType_USAGE_COST_TYPE_UNSPECIFIED: + return UsageCostTypeUnspecified + case pb.UsageCostType_USAGE_COST_TYPE_DATA_UPLOAD: + return UsageCostTypeDataUpload + case pb.UsageCostType_USAGE_COST_TYPE_DATA_EGRESS: + return UsageCostTypeDataEgress + case pb.UsageCostType_USAGE_COST_TYPE_REMOTE_CONTROL: + return UsageCostTypeRemoteControl + case pb.UsageCostType_USAGE_COST_TYPE_STANDARD_COMPUTE: + return UsageCostTypeStandardCompute + case pb.UsageCostType_USAGE_COST_TYPE_CLOUD_STORAGE: + return UsageCostTypeCloudStorage + case pb.UsageCostType_USAGE_COST_TYPE_BINARY_DATA_CLOUD_STORAGE: + return UsageCostTypeBinaryDataCloudStorage + case pb.UsageCostType_USAGE_COST_TYPE_OTHER_CLOUD_STORAGE: + return UsageCostTypeOtherCloudStorage + case pb.UsageCostType_USAGE_COST_TYPE_PER_MACHINE: + return UsageCostTypePerMachine + default: + return UsageCostTypeUnspecified + } +} + +type UsageCost struct { + ResourceType UsageCostType + Cost float64 +} + +func usageCostFromProto(cost *pb.UsageCost) *UsageCost { + return &UsageCost{ + ResourceType: usageCostTypeFromProto(cost.ResourceType), + Cost: cost.Cost, + } +} + +type ResourceUsageCosts struct { + UsageCosts []*UsageCost + Discount float64 + TotalWithDiscount float64 + TotalWithoutDiscount float64 +} + +func resourceUsageCostsFromProto(costs *pb.ResourceUsageCosts) *ResourceUsageCosts { + var usageCosts []*UsageCost + for _, cost := range(costs.UsageCosts) { + usageCosts = append(usageCosts, usageCostFromProto(cost)) + } + return &ResourceUsageCosts{ + UsageCosts: usageCosts, + Discount: costs.Discount, + TotalWithDiscount: costs.TotalWithDiscount, + TotalWithoutDiscount: costs.TotalWithoutDiscount, + } +} + +type ResourceUsageCostsBySource struct { + SourceType pb.SourceType + ResourceUsageCosts *ResourceUsageCosts + TierName string +} + +func resourceUsageCostsBySourceFromProto(costs *pb.ResourceUsageCostsBySource) *ResourceUsageCostsBySource { + return &ResourceUsageCostsBySource{ + SourceType: costs.SourceType, + ResourceUsageCosts: resourceUsageCostsFromProto(costs.ResourceUsageCosts), + TierName: costs.TierName, + } +} + +type GetCurrentMonthUsageResponse struct { + StartDate *timestamppb.Timestamp + EndDate *timestamppb.Timestamp + ResourceUsageCostsBySource []*ResourceUsageCostsBySource + Subtotal float64 +} + +func getCurrentMonthUsageResponseFromProto(response *pb.GetCurrentMonthUsageResponse) *GetCurrentMonthUsageResponse { + var costs []*ResourceUsageCostsBySource + for _, cost := range(response.ResourceUsageCostsBySource) { + costs = append(costs, resourceUsageCostsBySourceFromProto(cost)) + } + return &GetCurrentMonthUsageResponse{ + StartDate: response.StartDate, + EndDate: response.EndDate, + ResourceUsageCostsBySource: costs, + Subtotal: response.Subtotal, + } +} + +type PaymentMethodType int32 + +const ( + PaymentMethodTypeUnspecified PaymentMethodType = iota + PaymentMethodtypeCard +) + +func paymentMethodTypeFromProto(methodType pb.PaymentMethodType) PaymentMethodType { + switch methodType { + case pb.PaymentMethodType_PAYMENT_METHOD_TYPE_UNSPECIFIED: + return PaymentMethodTypeUnspecified + case pb.PaymentMethodType_PAYMENT_METHOD_TYPE_CARD: + return PaymentMethodtypeCard + default: + return PaymentMethodTypeUnspecified + } +} + +type PaymentMethodCard struct { + Brand string + LastFourDigits string +} + +func paymentMethodCardFromProto(card *pb.PaymentMethodCard) *PaymentMethodCard { + return &PaymentMethodCard{ + Brand: card.Brand, + LastFourDigits: card.LastFourDigits, + } +} + +type GetOrgBillingInformationResponse struct { + Type PaymentMethodType + BillingEmail string + // defined if type is PaymentMethodTypeCard + Method *PaymentMethodCard + // only return for billing dashboard admin users + BillingTier *string +} + +func getOrgBillingInformationResponseFromProto(resp *pb.GetOrgBillingInformationResponse) *GetOrgBillingInformationResponse { + return &GetOrgBillingInformationResponse{ + Type: paymentMethodTypeFromProto(resp.Type), + BillingEmail: resp.BillingEmail, + Method: paymentMethodCardFromProto(resp.Method), + BillingTier: resp.BillingTier, + } +} + +type InvoiceSummary struct { + ID string + InvoiceDate *timestamppb.Timestamp + InvoiceAmount float64 + Status string + DueDate *timestamppb.Timestamp + PaidDate *timestamppb.Timestamp +} + +func invoiceSummaryFromProto(summary *pb.InvoiceSummary) *InvoiceSummary { + return &InvoiceSummary{ + ID: summary.Id, + InvoiceDate: summary.InvoiceDate, + InvoiceAmount: summary.InvoiceAmount, + Status: summary.Status, + DueDate: summary.DueDate, + PaidDate: summary.PaidDate, + } +} + +type BillingClient struct { + client pb.BillingServiceClient +} + +func NewBillingClient(conn rpc.ClientConn) *BillingClient { + return &BillingClient{client: pb.NewBillingServiceClient(conn)} +} + +func (c *BillingClient) GetCurrentMonthUsage(ctx context.Context, orgID string) (*GetCurrentMonthUsageResponse, error) { + resp, err := c.client.GetCurrentMonthUsage(ctx, &pb.GetCurrentMonthUsageRequest{ + OrgId: orgID, + }) + if err != nil { + return nil, err + } + return getCurrentMonthUsageResponseFromProto(resp), nil +} + +func (c *BillingClient) GetOrgBillingInformation(ctx context.Context, orgID string) (*GetOrgBillingInformationResponse, error) { + resp, err := c.client.GetOrgBillingInformation(ctx, &pb.GetOrgBillingInformationRequest{ + OrgId: orgID, + }) + if err != nil { + return nil, err + } + return getOrgBillingInformationResponseFromProto(resp), nil +} + +func (c *BillingClient) GetInvoicesSummary(ctx context.Context, orgID string) (float64, []*InvoiceSummary, error) { + resp, err := c.client.GetInvoicesSummary(ctx, &pb.GetInvoicesSummaryRequest{ + OrgId: orgID, + }) + if err != nil { + return 0, nil, err + } + var invoices []*InvoiceSummary + for _, invoice := range(resp.Invoices) { + invoices = append(invoices, invoiceSummaryFromProto(invoice)) + } + return resp.OutstandingBalance, invoices, nil +} + +func (c *BillingClient) GetInvoicePdf(ctx context.Context, id, orgID string) () {} + +func (c *BillingClient) SendPaymentRequiredEmail(ctx context.Context, customerOrgID, billingOwnerOrgID string) error { + _, err := c.client.SendPaymentRequiredEmail(ctx, &pb.SendPaymentRequiredEmailRequest{ + CustomerOrgId: customerOrgID, + BillingOwnerOrgId: billingOwnerOrgID, + }) + return err +} From 229c25727dfdcccc49c1c22eadc4778b5f1ea613 Mon Sep 17 00:00:00 2001 From: purplenicole730 Date: Wed, 20 Nov 2024 11:59:04 -0500 Subject: [PATCH 02/15] add comments --- app/app_client.go | 1126 +++++++++++++++++++++++++++++++++++++++++ app/billing_client.go | 56 +- 2 files changed, 1180 insertions(+), 2 deletions(-) create mode 100644 app/app_client.go diff --git a/app/app_client.go b/app/app_client.go new file mode 100644 index 00000000000..c2ebc2700df --- /dev/null +++ b/app/app_client.go @@ -0,0 +1,1126 @@ +// Package app contains the interfaces that manage a machine fleet with code instead of with the graphical interface of the Viam App. +// +// [fleet management docs]: https://docs.viam.com/appendix/apis/fleet/ +package app + +import ( + "context" + "sync" + + packages "go.viam.com/api/app/packages/v1" + pb "go.viam.com/api/app/v1" + "go.viam.com/utils/protoutils" + "go.viam.com/utils/rpc" + "google.golang.org/protobuf/types/known/timestamppb" + + "go.viam.com/rdk/logging" +) + +// AppClient is a gRPC client for method calls to the App API. +// +//nolint:revive // stutter: Ignore the "stuttering" warning for this type name +type AppClient struct { + client pb.AppServiceClient + logger logging.Logger + + mu sync.Mutex +} + +// NewAppClient constructs a new AppClient using the connection passed in by the Viam client and the provided logger. +func NewAppClient(conn rpc.ClientConn, logger logging.Logger) *AppClient { + return &AppClient{client: pb.NewAppServiceClient(conn), logger: logger} +} + +// GetUserIDByEmail gets the ID of the user with the given email. +func (c *AppClient) GetUserIDByEmail(ctx context.Context, email string) (string, error) { + resp, err := c.client.GetUserIDByEmail(ctx, &pb.GetUserIDByEmailRequest{ + Email: email, + }) + if err != nil { + return "", err + } + return resp.UserId, nil +} + +// CreateOrganization creates a new organization. +func (c *AppClient) CreateOrganization(ctx context.Context, name string) (*Organization, error) { + resp, err := c.client.CreateOrganization(ctx, &pb.CreateOrganizationRequest{ + Name: name, + }) + if err != nil { + return nil, err + } + return organizationFromProto(resp.Organization), nil +} + +// ListOrganizations lists all the organizations. +func (c *AppClient) ListOrganizations(ctx context.Context) ([]*Organization, error) { + resp, err := c.client.ListOrganizations(ctx, &pb.ListOrganizationsRequest{}) + if err != nil { + return nil, err + } + + var organizations []*Organization + for _, org := range resp.Organizations { + organizations = append(organizations, organizationFromProto(org)) + } + return organizations, nil +} + +// GetOrganizationsWithAccessToLocation gets all the organizations that have access to a location. +func (c *AppClient) GetOrganizationsWithAccessToLocation(ctx context.Context, locationID string) ([]*OrganizationIdentity, error) { + resp, err := c.client.GetOrganizationsWithAccessToLocation(ctx, &pb.GetOrganizationsWithAccessToLocationRequest{ + LocationId: locationID, + }) + if err != nil { + return nil, err + } + + var organizations []*OrganizationIdentity + for _, org := range resp.OrganizationIdentities { + organizations = append(organizations, organizationIdentityFromProto(org)) + } + return organizations, nil +} + +// ListOrganizationsByUser lists all the organizations that a user belongs to. +func (c *AppClient) ListOrganizationsByUser(ctx context.Context, userID string) ([]*OrgDetails, error) { + resp, err := c.client.ListOrganizationsByUser(ctx, &pb.ListOrganizationsByUserRequest{ + UserId: userID, + }) + if err != nil { + return nil, err + } + + var organizations []*OrgDetails + for _, org := range resp.Orgs { + organizations = append(organizations, orgDetailsFromProto(org)) + } + return organizations, nil +} + +// GetOrganization gets an organization. +func (c *AppClient) GetOrganization(ctx context.Context, orgID string) (*Organization, error) { + resp, err := c.client.GetOrganization(ctx, &pb.GetOrganizationRequest{ + OrganizationId: orgID, + }) + if err != nil { + return nil, err + } + return organizationFromProto(resp.Organization), nil +} + +// GetOrganizationNamespaceAvailability checks for namespace availability throughout all organizations. +func (c *AppClient) GetOrganizationNamespaceAvailability(ctx context.Context, namespace string) (bool, error) { + resp, err := c.client.GetOrganizationNamespaceAvailability(ctx, &pb.GetOrganizationNamespaceAvailabilityRequest{ + PublicNamespace: namespace, + }) + if err != nil { + return false, err + } + return resp.Available, nil +} + +// UpdateOrganization updates an organization. +func (c *AppClient) UpdateOrganization(ctx context.Context, orgID string, name, namespace, region, cid *string) (*Organization, error) { + resp, err := c.client.UpdateOrganization(ctx, &pb.UpdateOrganizationRequest{ + OrganizationId: orgID, + Name: name, + PublicNamespace: namespace, + Region: region, + Cid: cid, + }) + if err != nil { + return nil, err + } + return organizationFromProto(resp.Organization), nil +} + +// DeleteOrganization deletes an organization. +func (c *AppClient) DeleteOrganization(ctx context.Context, orgID string) error { + _, err := c.client.DeleteOrganization(ctx, &pb.DeleteOrganizationRequest{ + OrganizationId: orgID, + }) + return err +} + +// ListOrganizationMembers lists all members of an organization and all invited members to the organization. +func (c *AppClient) ListOrganizationMembers(ctx context.Context, orgID string) ([]*OrganizationMember, []*OrganizationInvite, error) { + resp, err := c.client.ListOrganizationMembers(ctx, &pb.ListOrganizationMembersRequest{ + OrganizationId: orgID, + }) + if err != nil { + return nil, nil, err + } + + var members []*OrganizationMember + for _, member := range resp.Members { + members = append(members, organizationMemberFromProto(member)) + } + var invites []*OrganizationInvite + for _, invite := range resp.Invites { + invites = append(invites, organizationInviteFromProto(invite)) + } + return members, invites, nil +} + +// CreateOrganizationInvite creates an organization invite to an organization. +func (c *AppClient) CreateOrganizationInvite( + ctx context.Context, orgID, email string, authorizations []*Authorization, sendEmailInvite *bool, +) (*OrganizationInvite, error) { + var pbAuthorizations []*pb.Authorization + for _, authorization := range authorizations { + pbAuthorizations = append(pbAuthorizations, authorizationToProto(authorization)) + } + resp, err := c.client.CreateOrganizationInvite(ctx, &pb.CreateOrganizationInviteRequest{ + OrganizationId: orgID, + Email: email, + Authorizations: pbAuthorizations, + SendEmailInvite: sendEmailInvite, + }) + if err != nil { + return nil, err + } + return organizationInviteFromProto(resp.Invite), nil +} + +// UpdateOrganizationInviteAuthorizations updates the authorizations attached to an organization invite. +func (c *AppClient) UpdateOrganizationInviteAuthorizations( + ctx context.Context, orgID, email string, addAuthorizations, removeAuthorizations []*Authorization, +) (*OrganizationInvite, error) { + var pbAddAuthorizations []*pb.Authorization + for _, authorization := range addAuthorizations { + pbAddAuthorizations = append(pbAddAuthorizations, authorizationToProto(authorization)) + } + var pbRemoveAuthorizations []*pb.Authorization + for _, authorization := range removeAuthorizations { + pbRemoveAuthorizations = append(pbRemoveAuthorizations, authorizationToProto(authorization)) + } + resp, err := c.client.UpdateOrganizationInviteAuthorizations(ctx, &pb.UpdateOrganizationInviteAuthorizationsRequest{ + OrganizationId: orgID, + Email: email, + AddAuthorizations: pbAddAuthorizations, + RemoveAuthorizations: pbRemoveAuthorizations, + }) + if err != nil { + return nil, err + } + return organizationInviteFromProto(resp.Invite), nil +} + +// DeleteOrganizationMember deletes an organization member from an organization. +func (c *AppClient) DeleteOrganizationMember(ctx context.Context, orgID, userID string) error { + _, err := c.client.DeleteOrganizationMember(ctx, &pb.DeleteOrganizationMemberRequest{ + OrganizationId: orgID, + UserId: userID, + }) + return err +} + +// DeleteOrganizationInvite deletes an organization invite. +func (c *AppClient) DeleteOrganizationInvite(ctx context.Context, orgID, email string) error { + _, err := c.client.DeleteOrganizationInvite(ctx, &pb.DeleteOrganizationInviteRequest{ + OrganizationId: orgID, + Email: email, + }) + return err +} + +// ResendOrganizationInvite resends an organization invite. +func (c *AppClient) ResendOrganizationInvite(ctx context.Context, orgID, email string) (*OrganizationInvite, error) { + resp, err := c.client.ResendOrganizationInvite(ctx, &pb.ResendOrganizationInviteRequest{ + OrganizationId: orgID, + Email: email, + }) + if err != nil { + return nil, err + } + return organizationInviteFromProto(resp.Invite), nil +} + +// EnableBillingService enables a billing service to an address in an organization. +func (c *AppClient) EnableBillingService(ctx context.Context, orgID string, billingAddress *BillingAddress) error { + _, err := c.client.EnableBillingService(ctx, &pb.EnableBillingServiceRequest{ + OrgId: orgID, + BillingAddress: billingAddressToProto(billingAddress), + }) + return err +} + +// DisableBillingService disables the billing service for an organization. +func (c *AppClient) DisableBillingService(ctx context.Context, orgID string) error { + _, err := c.client.DisableBillingService(ctx, &pb.DisableBillingServiceRequest{ + OrgId: orgID, + }) + return err +} + +// UpdateBillingService updates the billing service of an organization. +func (c *AppClient) UpdateBillingService( + ctx context.Context, orgID string, billingAddress *BillingAddress, billingSupportEmail string, +) error { + _, err := c.client.UpdateBillingService(ctx, &pb.UpdateBillingServiceRequest{ + OrgId: orgID, + BillingAddress: billingAddressToProto(billingAddress), + BillingSupportEmail: billingSupportEmail, + }) + return err +} + +// OrganizationSetSupportEmail sets an organization's support email. +func (c *AppClient) OrganizationSetSupportEmail(ctx context.Context, orgID, email string) error { + _, err := c.client.OrganizationSetSupportEmail(ctx, &pb.OrganizationSetSupportEmailRequest{ + OrgId: orgID, + Email: email, + }) + return err +} + +// OrganizationGetSupportEmail gets an organization's support email. +func (c *AppClient) OrganizationGetSupportEmail(ctx context.Context, orgID string) (string, error) { + resp, err := c.client.OrganizationGetSupportEmail(ctx, &pb.OrganizationGetSupportEmailRequest{ + OrgId: orgID, + }) + if err != nil { + return "", err + } + return resp.Email, nil +} + +// CreateLocation creates a location. +func (c *AppClient) CreateLocation(ctx context.Context, orgID, name string, parentLocationID *string) (*Location, error) { + resp, err := c.client.CreateLocation(ctx, &pb.CreateLocationRequest{ + OrganizationId: orgID, + Name: name, + ParentLocationId: parentLocationID, + }) + if err != nil { + return nil, err + } + return locationFromProto(resp.Location), nil +} + +// GetLocation gets a location. +func (c *AppClient) GetLocation(ctx context.Context, locationID string) (*Location, error) { + resp, err := c.client.GetLocation(ctx, &pb.GetLocationRequest{ + LocationId: locationID, + }) + if err != nil { + return nil, err + } + return locationFromProto(resp.Location), nil +} + +// UpdateLocation updates a location. +func (c *AppClient) UpdateLocation(ctx context.Context, locationID string, name, parentLocationID, region *string) (*Location, error) { + resp, err := c.client.UpdateLocation(ctx, &pb.UpdateLocationRequest{ + LocationId: locationID, + Name: name, + ParentLocationId: parentLocationID, + Region: region, + }) + if err != nil { + return nil, err + } + return locationFromProto(resp.Location), nil +} + +// DeleteLocation deletes a location. +func (c *AppClient) DeleteLocation(ctx context.Context, locationID string) error { + _, err := c.client.DeleteLocation(ctx, &pb.DeleteLocationRequest{ + LocationId: locationID, + }) + return err +} + +// ListLocations gets a list of locations under the specified organization. +func (c *AppClient) ListLocations(ctx context.Context, orgID string) ([]*Location, error) { + resp, err := c.client.ListLocations(ctx, &pb.ListLocationsRequest{ + OrganizationId: orgID, + }) + if err != nil { + return nil, err + } + + var locations []*Location + for _, location := range resp.Locations { + locations = append(locations, locationFromProto(location)) + } + return locations, nil +} + +// ShareLocation shares a location with an organization. +func (c *AppClient) ShareLocation(ctx context.Context, locationID, orgID string) error { + _, err := c.client.ShareLocation(ctx, &pb.ShareLocationRequest{ + LocationId: locationID, + OrganizationId: orgID, + }) + return err +} + +// UnshareLocation stops sharing a location with an organization. +func (c *AppClient) UnshareLocation(ctx context.Context, locationID, orgID string) error { + _, err := c.client.UnshareLocation(ctx, &pb.UnshareLocationRequest{ + LocationId: locationID, + OrganizationId: orgID, + }) + return err +} + +// LocationAuth gets a location's authorization secrets. +func (c *AppClient) LocationAuth(ctx context.Context, locationID string) (*LocationAuth, error) { + resp, err := c.client.LocationAuth(ctx, &pb.LocationAuthRequest{ + LocationId: locationID, + }) + if err != nil { + return nil, err + } + return locationAuthFromProto(resp.Auth), nil +} + +// CreateLocationSecret creates a new generated secret in the location. Succeeds if there are no more than 2 active secrets after creation. +func (c *AppClient) CreateLocationSecret(ctx context.Context, locationID string) (*LocationAuth, error) { + resp, err := c.client.CreateLocationSecret(ctx, &pb.CreateLocationSecretRequest{ + LocationId: locationID, + }) + if err != nil { + return nil, err + } + return locationAuthFromProto(resp.Auth), nil +} + +// DeleteLocationSecret deletes a secret from the location. +func (c *AppClient) DeleteLocationSecret(ctx context.Context, locationID, secretID string) error { + _, err := c.client.DeleteLocationSecret(ctx, &pb.DeleteLocationSecretRequest{ + LocationId: locationID, + SecretId: secretID, + }) + return err +} + +// GetRobot gets a specific robot by ID. +func (c *AppClient) GetRobot(ctx context.Context, id string) (*Robot, error) { + resp, err := c.client.GetRobot(ctx, &pb.GetRobotRequest{ + Id: id, + }) + if err != nil { + return nil, err + } + return robotFromProto(resp.Robot), nil +} + +// GetRoverRentalRobots gets rover rental robots within an organization. +func (c *AppClient) GetRoverRentalRobots(ctx context.Context, orgID string) ([]*RoverRentalRobot, error) { + resp, err := c.client.GetRoverRentalRobots(ctx, &pb.GetRoverRentalRobotsRequest{ + OrgId: orgID, + }) + if err != nil { + return nil, err + } + var robots []*RoverRentalRobot + for _, robot := range resp.Robots { + robots = append(robots, roverRentalRobotFromProto(robot)) + } + return robots, nil +} + +// GetRobotParts gets a list of all the parts under a specific machine. +func (c *AppClient) GetRobotParts(ctx context.Context, robotID string) ([]*RobotPart, error) { + resp, err := c.client.GetRobotParts(ctx, &pb.GetRobotPartsRequest{ + RobotId: robotID, + }) + if err != nil { + return nil, err + } + var parts []*RobotPart + for _, part := range resp.Parts { + parts = append(parts, robotPartFromProto(part)) + } + return parts, nil +} + +// GetRobotPart gets a specific robot part and its config by ID. +func (c *AppClient) GetRobotPart(ctx context.Context, id string) (*RobotPart, string, error) { + resp, err := c.client.GetRobotPart(ctx, &pb.GetRobotPartRequest{ + Id: id, + }) + if err != nil { + return nil, "", err + } + return robotPartFromProto(resp.Part), resp.ConfigJson, nil +} + +// GetRobotPartLogs gets the logs associated with a robot part and the next page token, +// defaulting to the most recent page if pageToken is empty. Logs of all levels are returned when levels is empty. +func (c *AppClient) GetRobotPartLogs( + ctx context.Context, + id string, + filter, + pageToken *string, + levels []string, + start, + end *timestamppb.Timestamp, + limit *int64, + source *string, +) ([]*LogEntry, string, error) { + resp, err := c.client.GetRobotPartLogs(ctx, &pb.GetRobotPartLogsRequest{ + Id: id, + Filter: filter, + PageToken: pageToken, + Levels: levels, + Start: start, + End: end, + Limit: limit, + Source: source, + }) + if err != nil { + return nil, "", err + } + var logs []*LogEntry + for _, log := range resp.Logs { + logs = append(logs, logEntryFromProto(log)) + } + return logs, resp.NextPageToken, nil +} + +// TailRobotPartLogs gets a stream of log entries for a specific robot part. Logs are ordered by newest first. +func (c *AppClient) TailRobotPartLogs(ctx context.Context, id string, errorsOnly bool, filter *string, ch chan []*LogEntry) error { + stream := &robotPartLogStream{client: c} + + err := stream.startStream(ctx, id, errorsOnly, filter, ch) + if err != nil { + return err + } + + c.mu.Lock() + defer c.mu.Unlock() + return nil +} + +// GetRobotPartHistory gets a specific robot part history by ID. +func (c *AppClient) GetRobotPartHistory(ctx context.Context, id string) ([]*RobotPartHistoryEntry, error) { + resp, err := c.client.GetRobotPartHistory(ctx, &pb.GetRobotPartHistoryRequest{ + Id: id, + }) + if err != nil { + return nil, err + } + var history []*RobotPartHistoryEntry + for _, entry := range resp.History { + history = append(history, robotPartHistoryEntryFromProto(entry)) + } + return history, nil +} + +// UpdateRobotPart updates a robot part. +func (c *AppClient) UpdateRobotPart(ctx context.Context, id, name string, robotConfig interface{}) (*RobotPart, error) { + config, err := protoutils.StructToStructPb(robotConfig) + if err != nil { + return nil, err + } + resp, err := c.client.UpdateRobotPart(ctx, &pb.UpdateRobotPartRequest{ + Id: id, + Name: name, + RobotConfig: config, + }) + if err != nil { + return nil, err + } + return robotPartFromProto(resp.Part), nil +} + +// NewRobotPart creates a new robot part and returns its ID. +func (c *AppClient) NewRobotPart(ctx context.Context, robotID, partName string) (string, error) { + resp, err := c.client.NewRobotPart(ctx, &pb.NewRobotPartRequest{ + RobotId: robotID, + PartName: partName, + }) + if err != nil { + return "", err + } + return resp.PartId, nil +} + +// DeleteRobotPart deletes a robot part. +func (c *AppClient) DeleteRobotPart(ctx context.Context, partID string) error { + _, err := c.client.DeleteRobotPart(ctx, &pb.DeleteRobotPartRequest{ + PartId: partID, + }) + return err +} + +// GetRobotAPIKeys gets the robot API keys for the robot. +func (c *AppClient) GetRobotAPIKeys(ctx context.Context, robotID string) ([]*APIKeyWithAuthorizations, error) { + resp, err := c.client.GetRobotAPIKeys(ctx, &pb.GetRobotAPIKeysRequest{ + RobotId: robotID, + }) + if err != nil { + return nil, err + } + var keys []*APIKeyWithAuthorizations + for _, key := range resp.ApiKeys { + keys = append(keys, apiKeyWithAuthorizationsFromProto(key)) + } + return keys, nil +} + +// MarkPartAsMain marks the given part as the main part, and all the others as not. +func (c *AppClient) MarkPartAsMain(ctx context.Context, partID string) error { + _, err := c.client.MarkPartAsMain(ctx, &pb.MarkPartAsMainRequest{ + PartId: partID, + }) + return err +} + +// MarkPartForRestart marks the given part for restart. +// Once the robot part checks-in with the app the flag is reset on the robot part. +// Calling this multiple times before a robot part checks-in has no effect. +func (c *AppClient) MarkPartForRestart(ctx context.Context, partID string) error { + _, err := c.client.MarkPartForRestart(ctx, &pb.MarkPartForRestartRequest{ + PartId: partID, + }) + return err +} + +// CreateRobotPartSecret creates a new generated secret in the robot part. +// Succeeds if there are no more than 2 active secrets after creation. +func (c *AppClient) CreateRobotPartSecret(ctx context.Context, partID string) (*RobotPart, error) { + resp, err := c.client.CreateRobotPartSecret(ctx, &pb.CreateRobotPartSecretRequest{ + PartId: partID, + }) + if err != nil { + return nil, err + } + return robotPartFromProto(resp.Part), nil +} + +// DeleteRobotPartSecret deletes a secret from the robot part. +func (c *AppClient) DeleteRobotPartSecret(ctx context.Context, partID, secretID string) error { + _, err := c.client.DeleteRobotPartSecret(ctx, &pb.DeleteRobotPartSecretRequest{ + PartId: partID, + SecretId: secretID, + }) + return err +} + +// ListRobots gets a list of robots under a location. +func (c *AppClient) ListRobots(ctx context.Context, locationID string) ([]*Robot, error) { + resp, err := c.client.ListRobots(ctx, &pb.ListRobotsRequest{ + LocationId: locationID, + }) + if err != nil { + return nil, err + } + var robots []*Robot + for _, robot := range resp.Robots { + robots = append(robots, robotFromProto(robot)) + } + return robots, nil +} + +// NewRobot creates a new robot and returns its ID. +func (c *AppClient) NewRobot(ctx context.Context, name, location string) (string, error) { + resp, err := c.client.NewRobot(ctx, &pb.NewRobotRequest{ + Name: name, + Location: location, + }) + if err != nil { + return "", err + } + return resp.Id, nil +} + +// UpdateRobot updates a robot. +func (c *AppClient) UpdateRobot(ctx context.Context, id, name, location string) (*Robot, error) { + resp, err := c.client.UpdateRobot(ctx, &pb.UpdateRobotRequest{ + Id: id, + Name: name, + Location: location, + }) + if err != nil { + return nil, err + } + return robotFromProto(resp.Robot), nil +} + +// DeleteRobot deletes a robot. +func (c *AppClient) DeleteRobot(ctx context.Context, id string) error { + _, err := c.client.DeleteRobot(ctx, &pb.DeleteRobotRequest{ + Id: id, + }) + return err +} + +// ListFragments gets a list of fragments. +func (c *AppClient) ListFragments( + ctx context.Context, orgID string, showPublic bool, fragmentVisibility []FragmentVisibility, +) ([]*Fragment, error) { + var visibilities []pb.FragmentVisibility + for _, visibility := range fragmentVisibility { + pbFragmentVisibility := fragmentVisibilityToProto(visibility) + visibilities = append(visibilities, pbFragmentVisibility) + } + resp, err := c.client.ListFragments(ctx, &pb.ListFragmentsRequest{ + OrganizationId: orgID, + ShowPublic: showPublic, + FragmentVisibility: visibilities, + }) + if err != nil { + return nil, err + } + var fragments []*Fragment + for _, fragment := range resp.Fragments { + fragments = append(fragments, fragmentFromProto(fragment)) + } + return fragments, nil +} + +// GetFragment gets a single fragment. +func (c *AppClient) GetFragment(ctx context.Context, id string) (*Fragment, error) { + resp, err := c.client.GetFragment(ctx, &pb.GetFragmentRequest{ + Id: id, + }) + if err != nil { + return nil, err + } + return fragmentFromProto(resp.Fragment), nil +} + +// CreateFragment creates a fragment. +func (c *AppClient) CreateFragment( + ctx context.Context, name string, config interface{}, orgID string, visibility *FragmentVisibility, +) (*Fragment, error) { + cfg, err := protoutils.StructToStructPb(config) + if err != nil { + return nil, err + } + pbFragmentVisibility := fragmentVisibilityToProto(*visibility) + resp, err := c.client.CreateFragment(ctx, &pb.CreateFragmentRequest{ + Name: name, + Config: cfg, + OrganizationId: orgID, + Visibility: &pbFragmentVisibility, + }) + if err != nil { + return nil, err + } + return fragmentFromProto(resp.Fragment), nil +} + +// UpdateFragment updates a fragment. +func (c *AppClient) UpdateFragment( + ctx context.Context, id, name string, config map[string]interface{}, public *bool, visibility *FragmentVisibility, +) (*Fragment, error) { + cfg, err := protoutils.StructToStructPb(config) + if err != nil { + return nil, err + } + pbVisibility := fragmentVisibilityToProto(*visibility) + resp, err := c.client.UpdateFragment(ctx, &pb.UpdateFragmentRequest{ + Id: id, + Name: name, + Config: cfg, + Public: public, + Visibility: &pbVisibility, + }) + if err != nil { + return nil, err + } + return fragmentFromProto(resp.Fragment), nil +} + +// DeleteFragment deletes a fragment. +func (c *AppClient) DeleteFragment(ctx context.Context, id string) error { + _, err := c.client.DeleteFragment(ctx, &pb.DeleteFragmentRequest{ + Id: id, + }) + return err +} + +// ListMachineFragments gets top level and nested fragments for a amchine, as well as any other fragments specified by IDs. Additional +// fragments are useful when needing to view fragments that will be provisionally added to the machine alongside existing fragments. +func (c *AppClient) ListMachineFragments(ctx context.Context, machineID string, additionalFragmentIDs []string) ([]*Fragment, error) { + resp, err := c.client.ListMachineFragments(ctx, &pb.ListMachineFragmentsRequest{ + MachineId: machineID, + AdditionalFragmentIds: additionalFragmentIDs, + }) + if err != nil { + return nil, err + } + var fragments []*Fragment + for _, fragment := range resp.Fragments { + fragments = append(fragments, fragmentFromProto(fragment)) + } + return fragments, nil +} + +// GetFragmentHistory gets the fragment's history and the next page token. +func (c *AppClient) GetFragmentHistory( + ctx context.Context, id string, pageToken *string, pageLimit *int64, +) ([]*FragmentHistoryEntry, string, error) { + resp, err := c.client.GetFragmentHistory(ctx, &pb.GetFragmentHistoryRequest{ + Id: id, + PageToken: pageToken, + PageLimit: pageLimit, + }) + if err != nil { + return nil, "", err + } + var history []*FragmentHistoryEntry + for _, entry := range resp.History { + history = append(history, fragmentHistoryEntryFromProto(entry)) + } + return history, resp.NextPageToken, nil +} + +// AddRole creates an identity authorization. +func (c *AppClient) AddRole(ctx context.Context, orgID, identityID, role, resourceType, resourceID string) error { + authorization, err := createAuthorization(orgID, identityID, "", role, resourceType, resourceID) + if err != nil { + return err + } + _, err = c.client.AddRole(ctx, &pb.AddRoleRequest{ + Authorization: authorization, + }) + return err +} + +// RemoveRole deletes an identity authorization. +func (c *AppClient) RemoveRole(ctx context.Context, orgID, identityID, role, resourceType, resourceID string) error { + authorization, err := createAuthorization(orgID, identityID, "", role, resourceType, resourceID) + if err != nil { + return err + } + _, err = c.client.RemoveRole(ctx, &pb.RemoveRoleRequest{ + Authorization: authorization, + }) + return err +} + +// ChangeRole changes an identity authorization to a new identity authorization. +func (c *AppClient) ChangeRole( + ctx context.Context, + oldOrgID, + oldIdentityID, + oldRole, + oldResourceType, + oldResourceID, + newOrgID, + newIdentityID, + newRole, + newResourceType, + newResourceID string, +) error { + oldAuthorization, err := createAuthorization(oldOrgID, oldIdentityID, "", oldRole, oldResourceType, oldResourceID) + if err != nil { + return err + } + newAuthorization, err := createAuthorization(newOrgID, newIdentityID, "", newRole, newResourceType, newResourceID) + if err != nil { + return err + } + _, err = c.client.ChangeRole(ctx, &pb.ChangeRoleRequest{ + OldAuthorization: oldAuthorization, + NewAuthorization: newAuthorization, + }) + return err +} + +// ListAuthorizations returns all authorization roles for any given resources. +// If no resources are given, all resources within the organization will be included. +func (c *AppClient) ListAuthorizations(ctx context.Context, orgID string, resourceIDs []string) ([]*Authorization, error) { + resp, err := c.client.ListAuthorizations(ctx, &pb.ListAuthorizationsRequest{ + OrganizationId: orgID, + ResourceIds: resourceIDs, + }) + if err != nil { + return nil, err + } + var authorizations []*Authorization + for _, authorization := range resp.Authorizations { + authorizations = append(authorizations, authorizationFromProto(authorization)) + } + return authorizations, nil +} + +// CheckPermissions checks the validity of a list of permissions. +func (c *AppClient) CheckPermissions(ctx context.Context, permissions []*AuthorizedPermissions) ([]*AuthorizedPermissions, error) { + var pbPermissions []*pb.AuthorizedPermissions + for _, permission := range permissions { + pbPermissions = append(pbPermissions, authorizedPermissionsToProto(permission)) + } + + resp, err := c.client.CheckPermissions(ctx, &pb.CheckPermissionsRequest{ + Permissions: pbPermissions, + }) + if err != nil { + return nil, err + } + + var authorizedPermissions []*AuthorizedPermissions + for _, permission := range resp.AuthorizedPermissions { + authorizedPermissions = append(authorizedPermissions, authorizedPermissionsFromProto(permission)) + } + return authorizedPermissions, nil +} + +// GetRegistryItem gets a registry item. +func (c *AppClient) GetRegistryItem(ctx context.Context, itemID string) (*RegistryItem, error) { + resp, err := c.client.GetRegistryItem(ctx, &pb.GetRegistryItemRequest{ + ItemId: itemID, + }) + if err != nil { + return nil, err + } + item, err := registryItemFromProto(resp.Item) + if err != nil { + return nil, err + } + return item, nil +} + +// CreateRegistryItem creates a registry item. +func (c *AppClient) CreateRegistryItem(ctx context.Context, orgID, name string, packageType PackageType) error { + _, err := c.client.CreateRegistryItem(ctx, &pb.CreateRegistryItemRequest{ + OrganizationId: orgID, + Name: name, + Type: packageTypeToProto(packageType), + }) + return err +} + +// UpdateRegistryItem updates a registry item. +func (c *AppClient) UpdateRegistryItem( + ctx context.Context, itemID string, packageType PackageType, description string, visibility Visibility, url *string, +) error { + _, err := c.client.UpdateRegistryItem(ctx, &pb.UpdateRegistryItemRequest{ + ItemId: itemID, + Type: packageTypeToProto(packageType), + Description: description, + Visibility: visibilityToProto(visibility), + Url: url, + }) + return err +} + +// ListRegistryItems lists the registry items in an organization. +func (c *AppClient) ListRegistryItems( + ctx context.Context, + orgID *string, + types []PackageType, + visibilities []Visibility, + platforms []string, + statuses []RegistryItemStatus, + searchTerm, + pageToken *string, + publicNamespaces []string, +) ([]*RegistryItem, error) { + var pbTypes []packages.PackageType + for _, packageType := range types { + pbTypes = append(pbTypes, packageTypeToProto(packageType)) + } + var pbVisibilities []pb.Visibility + for _, visibility := range visibilities { + pbVisibilities = append(pbVisibilities, visibilityToProto(visibility)) + } + var pbStatuses []pb.RegistryItemStatus + for _, status := range statuses { + pbStatuses = append(pbStatuses, registryItemStatusToProto(status)) + } + resp, err := c.client.ListRegistryItems(ctx, &pb.ListRegistryItemsRequest{ + OrganizationId: orgID, + Types: pbTypes, + Visibilities: pbVisibilities, + Platforms: platforms, + Statuses: pbStatuses, + SearchTerm: searchTerm, + PageToken: pageToken, + PublicNamespaces: publicNamespaces, + }) + if err != nil { + return nil, err + } + var items []*RegistryItem + for _, item := range resp.Items { + i, err := registryItemFromProto(item) + if err != nil { + return nil, err + } + items = append(items, i) + } + return items, nil +} + +// DeleteRegistryItem deletes a registry item given an ID that is formatted as `prefix:name“ +// where `prefix“ is the owner's organization ID or namespace. +func (c *AppClient) DeleteRegistryItem(ctx context.Context, itemID string) error { + _, err := c.client.DeleteRegistryItem(ctx, &pb.DeleteRegistryItemRequest{ + ItemId: itemID, + }) + return err +} + +// TransferRegistryItem transfers a registry item to a namespace. +func (c *AppClient) TransferRegistryItem(ctx context.Context, itemID, newPublicNamespace string) error { + _, err := c.client.TransferRegistryItem(ctx, &pb.TransferRegistryItemRequest{ + ItemId: itemID, + NewPublicNamespace: newPublicNamespace, + }) + return err +} + +// CreateModule creates a module and returns its ID and URL. +func (c *AppClient) CreateModule(ctx context.Context, orgID, name string) (string, string, error) { + resp, err := c.client.CreateModule(ctx, &pb.CreateModuleRequest{ + OrganizationId: orgID, + Name: name, + }) + if err != nil { + return "", "", err + } + return resp.ModuleId, resp.Url, nil +} + +// UpdateModule updates the documentation URL, description, models, entrypoint, and/or the visibility of a module and returns its URL. +// A path to a setup script can be added that is run before a newly downloaded module starts. +func (c *AppClient) UpdateModule( + ctx context.Context, moduleID string, visibility Visibility, url, description string, models []*Model, entrypoint string, firstRun *string, +) (string, error) { + var pbModels []*pb.Model + for _, model := range models { + pbModels = append(pbModels, modelToProto(model)) + } + resp, err := c.client.UpdateModule(ctx, &pb.UpdateModuleRequest{ + ModuleId: moduleID, + Visibility: visibilityToProto(visibility), + Url: url, + Description: description, + Models: pbModels, + Entrypoint: entrypoint, + FirstRun: firstRun, + }) + if err != nil { + return "", err + } + return resp.Url, nil +} + +// UploadModuleFile uploads a module file and returns the URL of the uploaded file. +func (c *AppClient) UploadModuleFile(ctx context.Context, fileInfo ModuleFileInfo, file []byte) (string, error) { + stream := &uploadModuleFileStream{client: c} + url, err := stream.startStream(ctx, &fileInfo, file) + if err != nil { + return "", err + } + + c.mu.Lock() + defer c.mu.Unlock() + return url, nil +} + +// GetModule gets a module. +func (c *AppClient) GetModule(ctx context.Context, moduleID string) (*Module, error) { + resp, err := c.client.GetModule(ctx, &pb.GetModuleRequest{ + ModuleId: moduleID, + }) + if err != nil { + return nil, err + } + return moduleFromProto(resp.Module), nil +} + +// ListModules lists the modules in the organization. +func (c *AppClient) ListModules(ctx context.Context, orgID *string) ([]*Module, error) { + resp, err := c.client.ListModules(ctx, &pb.ListModulesRequest{ + OrganizationId: orgID, + }) + if err != nil { + return nil, err + } + var modules []*Module + for _, module := range resp.Modules { + modules = append(modules, moduleFromProto(module)) + } + return modules, nil +} + +// CreateKey creates a new API key associated with a list of authorizations and returns its key and ID. +func (c *AppClient) CreateKey( + ctx context.Context, orgID string, keyAuthorizations []APIKeyAuthorization, name string, +) (string, string, error) { + var authorizations []*pb.Authorization + for _, keyAuthorization := range keyAuthorizations { + authorization, err := createAuthorization( + orgID, "", "api-key", keyAuthorization.role, keyAuthorization.resourceType, keyAuthorization.resourceID) + if err != nil { + return "", "", err + } + authorizations = append(authorizations, authorization) + } + + resp, err := c.client.CreateKey(ctx, &pb.CreateKeyRequest{ + Authorizations: authorizations, + Name: name, + }) + if err != nil { + return "", "", err + } + return resp.Key, resp.Id, nil +} + +// DeleteKey deletes an API key. +func (c *AppClient) DeleteKey(ctx context.Context, id string) error { + _, err := c.client.DeleteKey(ctx, &pb.DeleteKeyRequest{ + Id: id, + }) + return err +} + +// ListKeys lists all the keys for the organization. +func (c *AppClient) ListKeys(ctx context.Context, orgID string) ([]*APIKeyWithAuthorizations, error) { + resp, err := c.client.ListKeys(ctx, &pb.ListKeysRequest{ + OrgId: orgID, + }) + if err != nil { + return nil, err + } + var apiKeys []*APIKeyWithAuthorizations + for _, key := range resp.ApiKeys { + apiKeys = append(apiKeys, apiKeyWithAuthorizationsFromProto(key)) + } + return apiKeys, nil +} + +// RenameKey renames an API key and returns its ID and name. +func (c *AppClient) RenameKey(ctx context.Context, id, name string) (string, string, error) { + resp, err := c.client.RenameKey(ctx, &pb.RenameKeyRequest{ + Id: id, + Name: name, + }) + if err != nil { + return "", "", err + } + return resp.Id, resp.Name, nil +} + +// RotateKey rotates an API key and returns its ID and key. +func (c *AppClient) RotateKey(ctx context.Context, id string) (string, string, error) { + resp, err := c.client.RotateKey(ctx, &pb.RotateKeyRequest{ + Id: id, + }) + if err != nil { + return "", "", err + } + return resp.Id, resp.Key, nil +} + +// CreateKeyFromExistingKeyAuthorizations creates a new API key with an existing key's authorizations and returns its ID and key. +func (c *AppClient) CreateKeyFromExistingKeyAuthorizations(ctx context.Context, id string) (string, string, error) { + resp, err := c.client.CreateKeyFromExistingKeyAuthorizations(ctx, &pb.CreateKeyFromExistingKeyAuthorizationsRequest{ + Id: id, + }) + if err != nil { + return "", "", err + } + return resp.Id, resp.Key, nil +} diff --git a/app/billing_client.go b/app/billing_client.go index 529ff3e1948..3cc381ade83 100644 --- a/app/billing_client.go +++ b/app/billing_client.go @@ -8,17 +8,27 @@ import ( "google.golang.org/protobuf/types/known/timestamppb" ) +// UsageCostType specifies the type of usage cost. type UsageCostType int32 const ( + // UsageCostTypeUnspecified is an unspecified usage cost type. UsageCostTypeUnspecified UsageCostType = iota + // UsageCostTypeDataUpload represents the usage cost from data upload. UsageCostTypeDataUpload + // UsageCostTypeDataEgress represents the usage cost from data egress. UsageCostTypeDataEgress + // UsageCostTypeRemoteControl represents the usage cost from remote control. UsageCostTypeRemoteControl + // UsageCostTypeStandardCompute represents the usage cost from standard compute. UsageCostTypeStandardCompute + // UsageCostTypeCloudStorage represents the usage cost from cloud storage. UsageCostTypeCloudStorage + // UsageCostTypeBinaryDataCloudStorage represents the usage cost from binary data cloud storage. UsageCostTypeBinaryDataCloudStorage + // UsageCostTypeOtherCloudStorage represents the usage cost from other cloud storage. UsageCostTypeOtherCloudStorage + // UsageCostTypePerMachine represents the usage cost per machine. UsageCostTypePerMachine ) @@ -47,6 +57,7 @@ func usageCostTypeFromProto(costType pb.UsageCostType) UsageCostType { } } +// UsageCost contains the cost and cost type. type UsageCost struct { ResourceType UsageCostType Cost float64 @@ -59,6 +70,7 @@ func usageCostFromProto(cost *pb.UsageCost) *UsageCost { } } +// ResourceUsageCosts holds the usage costs with discount information. type ResourceUsageCosts struct { UsageCosts []*UsageCost Discount float64 @@ -79,20 +91,47 @@ func resourceUsageCostsFromProto(costs *pb.ResourceUsageCosts) *ResourceUsageCos } } +// SourceType is the type of source from which a cost is coming from. +type SourceType int32 + +const ( + // SourceTypeUnspecified represents an unspecified source type. + SourceTypeUnspecified SourceType = iota + // SourceTypeOrg represents an organization. + SourceTypeOrg + // SourceTypeFragment represents a fragment. + SourceTypeFragment +) + +func sourceTypeFromProto(sourceType pb.SourceType) SourceType { + switch sourceType { + case pb.SourceType_SOURCE_TYPE_UNSPECIFIED: + return SourceTypeUnspecified + case pb.SourceType_SOURCE_TYPE_ORG: + return SourceTypeOrg + case pb.SourceType_SOURCE_TYPE_FRAGMENT: + return SourceTypeFragment + default: + return SourceTypeUnspecified + } +} + +// ResourceUsageCostsBySource contains the resource usage costs of a source type. type ResourceUsageCostsBySource struct { - SourceType pb.SourceType + SourceType SourceType ResourceUsageCosts *ResourceUsageCosts TierName string } func resourceUsageCostsBySourceFromProto(costs *pb.ResourceUsageCostsBySource) *ResourceUsageCostsBySource { return &ResourceUsageCostsBySource{ - SourceType: costs.SourceType, + SourceType: sourceTypeFromProto(costs.SourceType), ResourceUsageCosts: resourceUsageCostsFromProto(costs.ResourceUsageCosts), TierName: costs.TierName, } } +// GetCurrentMonthUsageResponse contains the current month usage information. type GetCurrentMonthUsageResponse struct { StartDate *timestamppb.Timestamp EndDate *timestamppb.Timestamp @@ -113,10 +152,13 @@ func getCurrentMonthUsageResponseFromProto(response *pb.GetCurrentMonthUsageResp } } +// PaymentMethodType is the type of payment method. type PaymentMethodType int32 const ( + // PaymentMethodTypeUnspecified represents an unspecified payment method. PaymentMethodTypeUnspecified PaymentMethodType = iota + // PaymentMethodtypeCard represents a payment by card. PaymentMethodtypeCard ) @@ -131,6 +173,7 @@ func paymentMethodTypeFromProto(methodType pb.PaymentMethodType) PaymentMethodTy } } +// PaymentMethodCard holds the information of a card used for payment. type PaymentMethodCard struct { Brand string LastFourDigits string @@ -143,6 +186,7 @@ func paymentMethodCardFromProto(card *pb.PaymentMethodCard) *PaymentMethodCard { } } +// GetOrgBillingInformationResponse contains the information of an organization's billing information. type GetOrgBillingInformationResponse struct { Type PaymentMethodType BillingEmail string @@ -161,6 +205,7 @@ func getOrgBillingInformationResponseFromProto(resp *pb.GetOrgBillingInformation } } +// InvoiceSummary holds the information of an invoice summary. type InvoiceSummary struct { ID string InvoiceDate *timestamppb.Timestamp @@ -181,14 +226,17 @@ func invoiceSummaryFromProto(summary *pb.InvoiceSummary) *InvoiceSummary { } } +// BillingClient is a gRPC client for method calls to the Billing API. type BillingClient struct { client pb.BillingServiceClient } func NewBillingClient(conn rpc.ClientConn) *BillingClient { return &BillingClient{client: pb.NewBillingServiceClient(conn)} +// NewBillingClient constructs a new BillingClient using the connection passed in by the Viam client. } +// GetCurrentMonthUsage gets the data usage information for the current month for an organization. func (c *BillingClient) GetCurrentMonthUsage(ctx context.Context, orgID string) (*GetCurrentMonthUsageResponse, error) { resp, err := c.client.GetCurrentMonthUsage(ctx, &pb.GetCurrentMonthUsageRequest{ OrgId: orgID, @@ -199,6 +247,7 @@ func (c *BillingClient) GetCurrentMonthUsage(ctx context.Context, orgID string) return getCurrentMonthUsageResponseFromProto(resp), nil } +// GetOrgBillingInformation gets the billing information of an organization. func (c *BillingClient) GetOrgBillingInformation(ctx context.Context, orgID string) (*GetOrgBillingInformationResponse, error) { resp, err := c.client.GetOrgBillingInformation(ctx, &pb.GetOrgBillingInformationRequest{ OrgId: orgID, @@ -209,6 +258,7 @@ func (c *BillingClient) GetOrgBillingInformation(ctx context.Context, orgID stri return getOrgBillingInformationResponseFromProto(resp), nil } +// GetInvoicesSummary returns the outstanding balance and the invoice summaries of an organization. func (c *BillingClient) GetInvoicesSummary(ctx context.Context, orgID string) (float64, []*InvoiceSummary, error) { resp, err := c.client.GetInvoicesSummary(ctx, &pb.GetInvoicesSummaryRequest{ OrgId: orgID, @@ -224,7 +274,9 @@ func (c *BillingClient) GetInvoicesSummary(ctx context.Context, orgID string) (f } func (c *BillingClient) GetInvoicePdf(ctx context.Context, id, orgID string) () {} +// GetInvoicePDF gets the invoice PDF data. +// SendPaymentRequiredEmail sends an email about payment requirement. func (c *BillingClient) SendPaymentRequiredEmail(ctx context.Context, customerOrgID, billingOwnerOrgID string) error { _, err := c.client.SendPaymentRequiredEmail(ctx, &pb.SendPaymentRequiredEmailRequest{ CustomerOrgId: customerOrgID, From f798059f3fa9ab6ae8480850f7384a178d21949f Mon Sep 17 00:00:00 2001 From: purplenicole730 Date: Wed, 20 Nov 2024 12:07:05 -0500 Subject: [PATCH 03/15] add stream --- app/billing_client.go | 97 +++++++++++++++++++++++++++++++++++++++++-- 1 file changed, 94 insertions(+), 3 deletions(-) diff --git a/app/billing_client.go b/app/billing_client.go index 3cc381ade83..fe5a2ef9c0a 100644 --- a/app/billing_client.go +++ b/app/billing_client.go @@ -2,8 +2,11 @@ package app import ( "context" + "sync" pb "go.viam.com/api/app/v1" + "go.viam.com/rdk/logging" + "go.viam.com/utils" "go.viam.com/utils/rpc" "google.golang.org/protobuf/types/known/timestamppb" ) @@ -226,14 +229,91 @@ func invoiceSummaryFromProto(summary *pb.InvoiceSummary) *InvoiceSummary { } } +type invoiceStream struct { + client *BillingClient + streamCancel context.CancelFunc + streamMu sync.Mutex + + activeBackgroundWorkers sync.WaitGroup +} + +func (s *invoiceStream) startStream(ctx context.Context, id, orgID string, ch chan []byte) error { + s.streamMu.Lock() + defer s.streamMu.Unlock() + + if ctx.Err() != nil { + return ctx.Err() + } + + ctx, cancel := context.WithCancel(ctx) + s.streamCancel = cancel + + select { + case <- ctx.Done(): + return ctx.Err() + default: + } + + // This call won't return any errors it had until the client tries to receive. + //nolint:errcheck + stream, _ := s.client.client.GetInvoicePdf(ctx, &pb.GetInvoicePdfRequest{ + Id: id, + OrgId: orgID, + }) + _, err := stream.Recv() + if err != nil { + s.client.logger.CError(ctx, err) + return err + } + + // Create a background go routine to receive from the server stream. + // We rely on calling the Done function here rather than in close stream + // since managed go calls that function when the routine exits. + s.activeBackgroundWorkers.Add(1) + utils.ManagedGo(func() { + s.receiveFromStream(ctx, stream, ch) + }, + s.activeBackgroundWorkers.Done) + return nil +} + +func (s *invoiceStream) receiveFromStream(ctx context.Context, stream pb.BillingService_GetInvoicePdfClient, ch chan []byte) { + defer s.streamCancel() + + // repeatedly receive from the stream + for { + select { + case <-ctx.Done(): + s.client.logger.Debug(ctx.Err()) + return + default: + } + streamResp, err := stream.Recv() + if err != nil { + // only debug log the context canceled error + s.client.logger.Debug(err) + return + } + // If there is a response, send to the channel. + var pdf []byte + for _, data := range streamResp.Chunk { + pdf = append(pdf, data) + } + ch <- pdf + } +} + // BillingClient is a gRPC client for method calls to the Billing API. type BillingClient struct { client pb.BillingServiceClient + logger logging.Logger + + mu sync.Mutex } -func NewBillingClient(conn rpc.ClientConn) *BillingClient { - return &BillingClient{client: pb.NewBillingServiceClient(conn)} // NewBillingClient constructs a new BillingClient using the connection passed in by the Viam client. +func NewBillingClient(conn rpc.ClientConn, logger logging.Logger) *BillingClient { + return &BillingClient{client: pb.NewBillingServiceClient(conn), logger: logger} } // GetCurrentMonthUsage gets the data usage information for the current month for an organization. @@ -273,8 +353,19 @@ func (c *BillingClient) GetInvoicesSummary(ctx context.Context, orgID string) (f return resp.OutstandingBalance, invoices, nil } -func (c *BillingClient) GetInvoicePdf(ctx context.Context, id, orgID string) () {} // GetInvoicePDF gets the invoice PDF data. +func (c *BillingClient) GetInvoicePDF(ctx context.Context, id, orgID string, ch chan []byte) (error) { + stream := &invoiceStream{client: c} + + err := stream.startStream(ctx, id, orgID, ch) + if err != nil { + return nil + } + + c.mu.Lock() + defer c.mu.Unlock() + return nil +} // SendPaymentRequiredEmail sends an email about payment requirement. func (c *BillingClient) SendPaymentRequiredEmail(ctx context.Context, customerOrgID, billingOwnerOrgID string) error { From ff89a02798c716bb15a9e021fba50c5ded379aa1 Mon Sep 17 00:00:00 2001 From: purplenicole730 Date: Wed, 20 Nov 2024 14:01:00 -0500 Subject: [PATCH 04/15] delete this --- app/app_client.go | 1126 --------------------------------------------- 1 file changed, 1126 deletions(-) delete mode 100644 app/app_client.go diff --git a/app/app_client.go b/app/app_client.go deleted file mode 100644 index c2ebc2700df..00000000000 --- a/app/app_client.go +++ /dev/null @@ -1,1126 +0,0 @@ -// Package app contains the interfaces that manage a machine fleet with code instead of with the graphical interface of the Viam App. -// -// [fleet management docs]: https://docs.viam.com/appendix/apis/fleet/ -package app - -import ( - "context" - "sync" - - packages "go.viam.com/api/app/packages/v1" - pb "go.viam.com/api/app/v1" - "go.viam.com/utils/protoutils" - "go.viam.com/utils/rpc" - "google.golang.org/protobuf/types/known/timestamppb" - - "go.viam.com/rdk/logging" -) - -// AppClient is a gRPC client for method calls to the App API. -// -//nolint:revive // stutter: Ignore the "stuttering" warning for this type name -type AppClient struct { - client pb.AppServiceClient - logger logging.Logger - - mu sync.Mutex -} - -// NewAppClient constructs a new AppClient using the connection passed in by the Viam client and the provided logger. -func NewAppClient(conn rpc.ClientConn, logger logging.Logger) *AppClient { - return &AppClient{client: pb.NewAppServiceClient(conn), logger: logger} -} - -// GetUserIDByEmail gets the ID of the user with the given email. -func (c *AppClient) GetUserIDByEmail(ctx context.Context, email string) (string, error) { - resp, err := c.client.GetUserIDByEmail(ctx, &pb.GetUserIDByEmailRequest{ - Email: email, - }) - if err != nil { - return "", err - } - return resp.UserId, nil -} - -// CreateOrganization creates a new organization. -func (c *AppClient) CreateOrganization(ctx context.Context, name string) (*Organization, error) { - resp, err := c.client.CreateOrganization(ctx, &pb.CreateOrganizationRequest{ - Name: name, - }) - if err != nil { - return nil, err - } - return organizationFromProto(resp.Organization), nil -} - -// ListOrganizations lists all the organizations. -func (c *AppClient) ListOrganizations(ctx context.Context) ([]*Organization, error) { - resp, err := c.client.ListOrganizations(ctx, &pb.ListOrganizationsRequest{}) - if err != nil { - return nil, err - } - - var organizations []*Organization - for _, org := range resp.Organizations { - organizations = append(organizations, organizationFromProto(org)) - } - return organizations, nil -} - -// GetOrganizationsWithAccessToLocation gets all the organizations that have access to a location. -func (c *AppClient) GetOrganizationsWithAccessToLocation(ctx context.Context, locationID string) ([]*OrganizationIdentity, error) { - resp, err := c.client.GetOrganizationsWithAccessToLocation(ctx, &pb.GetOrganizationsWithAccessToLocationRequest{ - LocationId: locationID, - }) - if err != nil { - return nil, err - } - - var organizations []*OrganizationIdentity - for _, org := range resp.OrganizationIdentities { - organizations = append(organizations, organizationIdentityFromProto(org)) - } - return organizations, nil -} - -// ListOrganizationsByUser lists all the organizations that a user belongs to. -func (c *AppClient) ListOrganizationsByUser(ctx context.Context, userID string) ([]*OrgDetails, error) { - resp, err := c.client.ListOrganizationsByUser(ctx, &pb.ListOrganizationsByUserRequest{ - UserId: userID, - }) - if err != nil { - return nil, err - } - - var organizations []*OrgDetails - for _, org := range resp.Orgs { - organizations = append(organizations, orgDetailsFromProto(org)) - } - return organizations, nil -} - -// GetOrganization gets an organization. -func (c *AppClient) GetOrganization(ctx context.Context, orgID string) (*Organization, error) { - resp, err := c.client.GetOrganization(ctx, &pb.GetOrganizationRequest{ - OrganizationId: orgID, - }) - if err != nil { - return nil, err - } - return organizationFromProto(resp.Organization), nil -} - -// GetOrganizationNamespaceAvailability checks for namespace availability throughout all organizations. -func (c *AppClient) GetOrganizationNamespaceAvailability(ctx context.Context, namespace string) (bool, error) { - resp, err := c.client.GetOrganizationNamespaceAvailability(ctx, &pb.GetOrganizationNamespaceAvailabilityRequest{ - PublicNamespace: namespace, - }) - if err != nil { - return false, err - } - return resp.Available, nil -} - -// UpdateOrganization updates an organization. -func (c *AppClient) UpdateOrganization(ctx context.Context, orgID string, name, namespace, region, cid *string) (*Organization, error) { - resp, err := c.client.UpdateOrganization(ctx, &pb.UpdateOrganizationRequest{ - OrganizationId: orgID, - Name: name, - PublicNamespace: namespace, - Region: region, - Cid: cid, - }) - if err != nil { - return nil, err - } - return organizationFromProto(resp.Organization), nil -} - -// DeleteOrganization deletes an organization. -func (c *AppClient) DeleteOrganization(ctx context.Context, orgID string) error { - _, err := c.client.DeleteOrganization(ctx, &pb.DeleteOrganizationRequest{ - OrganizationId: orgID, - }) - return err -} - -// ListOrganizationMembers lists all members of an organization and all invited members to the organization. -func (c *AppClient) ListOrganizationMembers(ctx context.Context, orgID string) ([]*OrganizationMember, []*OrganizationInvite, error) { - resp, err := c.client.ListOrganizationMembers(ctx, &pb.ListOrganizationMembersRequest{ - OrganizationId: orgID, - }) - if err != nil { - return nil, nil, err - } - - var members []*OrganizationMember - for _, member := range resp.Members { - members = append(members, organizationMemberFromProto(member)) - } - var invites []*OrganizationInvite - for _, invite := range resp.Invites { - invites = append(invites, organizationInviteFromProto(invite)) - } - return members, invites, nil -} - -// CreateOrganizationInvite creates an organization invite to an organization. -func (c *AppClient) CreateOrganizationInvite( - ctx context.Context, orgID, email string, authorizations []*Authorization, sendEmailInvite *bool, -) (*OrganizationInvite, error) { - var pbAuthorizations []*pb.Authorization - for _, authorization := range authorizations { - pbAuthorizations = append(pbAuthorizations, authorizationToProto(authorization)) - } - resp, err := c.client.CreateOrganizationInvite(ctx, &pb.CreateOrganizationInviteRequest{ - OrganizationId: orgID, - Email: email, - Authorizations: pbAuthorizations, - SendEmailInvite: sendEmailInvite, - }) - if err != nil { - return nil, err - } - return organizationInviteFromProto(resp.Invite), nil -} - -// UpdateOrganizationInviteAuthorizations updates the authorizations attached to an organization invite. -func (c *AppClient) UpdateOrganizationInviteAuthorizations( - ctx context.Context, orgID, email string, addAuthorizations, removeAuthorizations []*Authorization, -) (*OrganizationInvite, error) { - var pbAddAuthorizations []*pb.Authorization - for _, authorization := range addAuthorizations { - pbAddAuthorizations = append(pbAddAuthorizations, authorizationToProto(authorization)) - } - var pbRemoveAuthorizations []*pb.Authorization - for _, authorization := range removeAuthorizations { - pbRemoveAuthorizations = append(pbRemoveAuthorizations, authorizationToProto(authorization)) - } - resp, err := c.client.UpdateOrganizationInviteAuthorizations(ctx, &pb.UpdateOrganizationInviteAuthorizationsRequest{ - OrganizationId: orgID, - Email: email, - AddAuthorizations: pbAddAuthorizations, - RemoveAuthorizations: pbRemoveAuthorizations, - }) - if err != nil { - return nil, err - } - return organizationInviteFromProto(resp.Invite), nil -} - -// DeleteOrganizationMember deletes an organization member from an organization. -func (c *AppClient) DeleteOrganizationMember(ctx context.Context, orgID, userID string) error { - _, err := c.client.DeleteOrganizationMember(ctx, &pb.DeleteOrganizationMemberRequest{ - OrganizationId: orgID, - UserId: userID, - }) - return err -} - -// DeleteOrganizationInvite deletes an organization invite. -func (c *AppClient) DeleteOrganizationInvite(ctx context.Context, orgID, email string) error { - _, err := c.client.DeleteOrganizationInvite(ctx, &pb.DeleteOrganizationInviteRequest{ - OrganizationId: orgID, - Email: email, - }) - return err -} - -// ResendOrganizationInvite resends an organization invite. -func (c *AppClient) ResendOrganizationInvite(ctx context.Context, orgID, email string) (*OrganizationInvite, error) { - resp, err := c.client.ResendOrganizationInvite(ctx, &pb.ResendOrganizationInviteRequest{ - OrganizationId: orgID, - Email: email, - }) - if err != nil { - return nil, err - } - return organizationInviteFromProto(resp.Invite), nil -} - -// EnableBillingService enables a billing service to an address in an organization. -func (c *AppClient) EnableBillingService(ctx context.Context, orgID string, billingAddress *BillingAddress) error { - _, err := c.client.EnableBillingService(ctx, &pb.EnableBillingServiceRequest{ - OrgId: orgID, - BillingAddress: billingAddressToProto(billingAddress), - }) - return err -} - -// DisableBillingService disables the billing service for an organization. -func (c *AppClient) DisableBillingService(ctx context.Context, orgID string) error { - _, err := c.client.DisableBillingService(ctx, &pb.DisableBillingServiceRequest{ - OrgId: orgID, - }) - return err -} - -// UpdateBillingService updates the billing service of an organization. -func (c *AppClient) UpdateBillingService( - ctx context.Context, orgID string, billingAddress *BillingAddress, billingSupportEmail string, -) error { - _, err := c.client.UpdateBillingService(ctx, &pb.UpdateBillingServiceRequest{ - OrgId: orgID, - BillingAddress: billingAddressToProto(billingAddress), - BillingSupportEmail: billingSupportEmail, - }) - return err -} - -// OrganizationSetSupportEmail sets an organization's support email. -func (c *AppClient) OrganizationSetSupportEmail(ctx context.Context, orgID, email string) error { - _, err := c.client.OrganizationSetSupportEmail(ctx, &pb.OrganizationSetSupportEmailRequest{ - OrgId: orgID, - Email: email, - }) - return err -} - -// OrganizationGetSupportEmail gets an organization's support email. -func (c *AppClient) OrganizationGetSupportEmail(ctx context.Context, orgID string) (string, error) { - resp, err := c.client.OrganizationGetSupportEmail(ctx, &pb.OrganizationGetSupportEmailRequest{ - OrgId: orgID, - }) - if err != nil { - return "", err - } - return resp.Email, nil -} - -// CreateLocation creates a location. -func (c *AppClient) CreateLocation(ctx context.Context, orgID, name string, parentLocationID *string) (*Location, error) { - resp, err := c.client.CreateLocation(ctx, &pb.CreateLocationRequest{ - OrganizationId: orgID, - Name: name, - ParentLocationId: parentLocationID, - }) - if err != nil { - return nil, err - } - return locationFromProto(resp.Location), nil -} - -// GetLocation gets a location. -func (c *AppClient) GetLocation(ctx context.Context, locationID string) (*Location, error) { - resp, err := c.client.GetLocation(ctx, &pb.GetLocationRequest{ - LocationId: locationID, - }) - if err != nil { - return nil, err - } - return locationFromProto(resp.Location), nil -} - -// UpdateLocation updates a location. -func (c *AppClient) UpdateLocation(ctx context.Context, locationID string, name, parentLocationID, region *string) (*Location, error) { - resp, err := c.client.UpdateLocation(ctx, &pb.UpdateLocationRequest{ - LocationId: locationID, - Name: name, - ParentLocationId: parentLocationID, - Region: region, - }) - if err != nil { - return nil, err - } - return locationFromProto(resp.Location), nil -} - -// DeleteLocation deletes a location. -func (c *AppClient) DeleteLocation(ctx context.Context, locationID string) error { - _, err := c.client.DeleteLocation(ctx, &pb.DeleteLocationRequest{ - LocationId: locationID, - }) - return err -} - -// ListLocations gets a list of locations under the specified organization. -func (c *AppClient) ListLocations(ctx context.Context, orgID string) ([]*Location, error) { - resp, err := c.client.ListLocations(ctx, &pb.ListLocationsRequest{ - OrganizationId: orgID, - }) - if err != nil { - return nil, err - } - - var locations []*Location - for _, location := range resp.Locations { - locations = append(locations, locationFromProto(location)) - } - return locations, nil -} - -// ShareLocation shares a location with an organization. -func (c *AppClient) ShareLocation(ctx context.Context, locationID, orgID string) error { - _, err := c.client.ShareLocation(ctx, &pb.ShareLocationRequest{ - LocationId: locationID, - OrganizationId: orgID, - }) - return err -} - -// UnshareLocation stops sharing a location with an organization. -func (c *AppClient) UnshareLocation(ctx context.Context, locationID, orgID string) error { - _, err := c.client.UnshareLocation(ctx, &pb.UnshareLocationRequest{ - LocationId: locationID, - OrganizationId: orgID, - }) - return err -} - -// LocationAuth gets a location's authorization secrets. -func (c *AppClient) LocationAuth(ctx context.Context, locationID string) (*LocationAuth, error) { - resp, err := c.client.LocationAuth(ctx, &pb.LocationAuthRequest{ - LocationId: locationID, - }) - if err != nil { - return nil, err - } - return locationAuthFromProto(resp.Auth), nil -} - -// CreateLocationSecret creates a new generated secret in the location. Succeeds if there are no more than 2 active secrets after creation. -func (c *AppClient) CreateLocationSecret(ctx context.Context, locationID string) (*LocationAuth, error) { - resp, err := c.client.CreateLocationSecret(ctx, &pb.CreateLocationSecretRequest{ - LocationId: locationID, - }) - if err != nil { - return nil, err - } - return locationAuthFromProto(resp.Auth), nil -} - -// DeleteLocationSecret deletes a secret from the location. -func (c *AppClient) DeleteLocationSecret(ctx context.Context, locationID, secretID string) error { - _, err := c.client.DeleteLocationSecret(ctx, &pb.DeleteLocationSecretRequest{ - LocationId: locationID, - SecretId: secretID, - }) - return err -} - -// GetRobot gets a specific robot by ID. -func (c *AppClient) GetRobot(ctx context.Context, id string) (*Robot, error) { - resp, err := c.client.GetRobot(ctx, &pb.GetRobotRequest{ - Id: id, - }) - if err != nil { - return nil, err - } - return robotFromProto(resp.Robot), nil -} - -// GetRoverRentalRobots gets rover rental robots within an organization. -func (c *AppClient) GetRoverRentalRobots(ctx context.Context, orgID string) ([]*RoverRentalRobot, error) { - resp, err := c.client.GetRoverRentalRobots(ctx, &pb.GetRoverRentalRobotsRequest{ - OrgId: orgID, - }) - if err != nil { - return nil, err - } - var robots []*RoverRentalRobot - for _, robot := range resp.Robots { - robots = append(robots, roverRentalRobotFromProto(robot)) - } - return robots, nil -} - -// GetRobotParts gets a list of all the parts under a specific machine. -func (c *AppClient) GetRobotParts(ctx context.Context, robotID string) ([]*RobotPart, error) { - resp, err := c.client.GetRobotParts(ctx, &pb.GetRobotPartsRequest{ - RobotId: robotID, - }) - if err != nil { - return nil, err - } - var parts []*RobotPart - for _, part := range resp.Parts { - parts = append(parts, robotPartFromProto(part)) - } - return parts, nil -} - -// GetRobotPart gets a specific robot part and its config by ID. -func (c *AppClient) GetRobotPart(ctx context.Context, id string) (*RobotPart, string, error) { - resp, err := c.client.GetRobotPart(ctx, &pb.GetRobotPartRequest{ - Id: id, - }) - if err != nil { - return nil, "", err - } - return robotPartFromProto(resp.Part), resp.ConfigJson, nil -} - -// GetRobotPartLogs gets the logs associated with a robot part and the next page token, -// defaulting to the most recent page if pageToken is empty. Logs of all levels are returned when levels is empty. -func (c *AppClient) GetRobotPartLogs( - ctx context.Context, - id string, - filter, - pageToken *string, - levels []string, - start, - end *timestamppb.Timestamp, - limit *int64, - source *string, -) ([]*LogEntry, string, error) { - resp, err := c.client.GetRobotPartLogs(ctx, &pb.GetRobotPartLogsRequest{ - Id: id, - Filter: filter, - PageToken: pageToken, - Levels: levels, - Start: start, - End: end, - Limit: limit, - Source: source, - }) - if err != nil { - return nil, "", err - } - var logs []*LogEntry - for _, log := range resp.Logs { - logs = append(logs, logEntryFromProto(log)) - } - return logs, resp.NextPageToken, nil -} - -// TailRobotPartLogs gets a stream of log entries for a specific robot part. Logs are ordered by newest first. -func (c *AppClient) TailRobotPartLogs(ctx context.Context, id string, errorsOnly bool, filter *string, ch chan []*LogEntry) error { - stream := &robotPartLogStream{client: c} - - err := stream.startStream(ctx, id, errorsOnly, filter, ch) - if err != nil { - return err - } - - c.mu.Lock() - defer c.mu.Unlock() - return nil -} - -// GetRobotPartHistory gets a specific robot part history by ID. -func (c *AppClient) GetRobotPartHistory(ctx context.Context, id string) ([]*RobotPartHistoryEntry, error) { - resp, err := c.client.GetRobotPartHistory(ctx, &pb.GetRobotPartHistoryRequest{ - Id: id, - }) - if err != nil { - return nil, err - } - var history []*RobotPartHistoryEntry - for _, entry := range resp.History { - history = append(history, robotPartHistoryEntryFromProto(entry)) - } - return history, nil -} - -// UpdateRobotPart updates a robot part. -func (c *AppClient) UpdateRobotPart(ctx context.Context, id, name string, robotConfig interface{}) (*RobotPart, error) { - config, err := protoutils.StructToStructPb(robotConfig) - if err != nil { - return nil, err - } - resp, err := c.client.UpdateRobotPart(ctx, &pb.UpdateRobotPartRequest{ - Id: id, - Name: name, - RobotConfig: config, - }) - if err != nil { - return nil, err - } - return robotPartFromProto(resp.Part), nil -} - -// NewRobotPart creates a new robot part and returns its ID. -func (c *AppClient) NewRobotPart(ctx context.Context, robotID, partName string) (string, error) { - resp, err := c.client.NewRobotPart(ctx, &pb.NewRobotPartRequest{ - RobotId: robotID, - PartName: partName, - }) - if err != nil { - return "", err - } - return resp.PartId, nil -} - -// DeleteRobotPart deletes a robot part. -func (c *AppClient) DeleteRobotPart(ctx context.Context, partID string) error { - _, err := c.client.DeleteRobotPart(ctx, &pb.DeleteRobotPartRequest{ - PartId: partID, - }) - return err -} - -// GetRobotAPIKeys gets the robot API keys for the robot. -func (c *AppClient) GetRobotAPIKeys(ctx context.Context, robotID string) ([]*APIKeyWithAuthorizations, error) { - resp, err := c.client.GetRobotAPIKeys(ctx, &pb.GetRobotAPIKeysRequest{ - RobotId: robotID, - }) - if err != nil { - return nil, err - } - var keys []*APIKeyWithAuthorizations - for _, key := range resp.ApiKeys { - keys = append(keys, apiKeyWithAuthorizationsFromProto(key)) - } - return keys, nil -} - -// MarkPartAsMain marks the given part as the main part, and all the others as not. -func (c *AppClient) MarkPartAsMain(ctx context.Context, partID string) error { - _, err := c.client.MarkPartAsMain(ctx, &pb.MarkPartAsMainRequest{ - PartId: partID, - }) - return err -} - -// MarkPartForRestart marks the given part for restart. -// Once the robot part checks-in with the app the flag is reset on the robot part. -// Calling this multiple times before a robot part checks-in has no effect. -func (c *AppClient) MarkPartForRestart(ctx context.Context, partID string) error { - _, err := c.client.MarkPartForRestart(ctx, &pb.MarkPartForRestartRequest{ - PartId: partID, - }) - return err -} - -// CreateRobotPartSecret creates a new generated secret in the robot part. -// Succeeds if there are no more than 2 active secrets after creation. -func (c *AppClient) CreateRobotPartSecret(ctx context.Context, partID string) (*RobotPart, error) { - resp, err := c.client.CreateRobotPartSecret(ctx, &pb.CreateRobotPartSecretRequest{ - PartId: partID, - }) - if err != nil { - return nil, err - } - return robotPartFromProto(resp.Part), nil -} - -// DeleteRobotPartSecret deletes a secret from the robot part. -func (c *AppClient) DeleteRobotPartSecret(ctx context.Context, partID, secretID string) error { - _, err := c.client.DeleteRobotPartSecret(ctx, &pb.DeleteRobotPartSecretRequest{ - PartId: partID, - SecretId: secretID, - }) - return err -} - -// ListRobots gets a list of robots under a location. -func (c *AppClient) ListRobots(ctx context.Context, locationID string) ([]*Robot, error) { - resp, err := c.client.ListRobots(ctx, &pb.ListRobotsRequest{ - LocationId: locationID, - }) - if err != nil { - return nil, err - } - var robots []*Robot - for _, robot := range resp.Robots { - robots = append(robots, robotFromProto(robot)) - } - return robots, nil -} - -// NewRobot creates a new robot and returns its ID. -func (c *AppClient) NewRobot(ctx context.Context, name, location string) (string, error) { - resp, err := c.client.NewRobot(ctx, &pb.NewRobotRequest{ - Name: name, - Location: location, - }) - if err != nil { - return "", err - } - return resp.Id, nil -} - -// UpdateRobot updates a robot. -func (c *AppClient) UpdateRobot(ctx context.Context, id, name, location string) (*Robot, error) { - resp, err := c.client.UpdateRobot(ctx, &pb.UpdateRobotRequest{ - Id: id, - Name: name, - Location: location, - }) - if err != nil { - return nil, err - } - return robotFromProto(resp.Robot), nil -} - -// DeleteRobot deletes a robot. -func (c *AppClient) DeleteRobot(ctx context.Context, id string) error { - _, err := c.client.DeleteRobot(ctx, &pb.DeleteRobotRequest{ - Id: id, - }) - return err -} - -// ListFragments gets a list of fragments. -func (c *AppClient) ListFragments( - ctx context.Context, orgID string, showPublic bool, fragmentVisibility []FragmentVisibility, -) ([]*Fragment, error) { - var visibilities []pb.FragmentVisibility - for _, visibility := range fragmentVisibility { - pbFragmentVisibility := fragmentVisibilityToProto(visibility) - visibilities = append(visibilities, pbFragmentVisibility) - } - resp, err := c.client.ListFragments(ctx, &pb.ListFragmentsRequest{ - OrganizationId: orgID, - ShowPublic: showPublic, - FragmentVisibility: visibilities, - }) - if err != nil { - return nil, err - } - var fragments []*Fragment - for _, fragment := range resp.Fragments { - fragments = append(fragments, fragmentFromProto(fragment)) - } - return fragments, nil -} - -// GetFragment gets a single fragment. -func (c *AppClient) GetFragment(ctx context.Context, id string) (*Fragment, error) { - resp, err := c.client.GetFragment(ctx, &pb.GetFragmentRequest{ - Id: id, - }) - if err != nil { - return nil, err - } - return fragmentFromProto(resp.Fragment), nil -} - -// CreateFragment creates a fragment. -func (c *AppClient) CreateFragment( - ctx context.Context, name string, config interface{}, orgID string, visibility *FragmentVisibility, -) (*Fragment, error) { - cfg, err := protoutils.StructToStructPb(config) - if err != nil { - return nil, err - } - pbFragmentVisibility := fragmentVisibilityToProto(*visibility) - resp, err := c.client.CreateFragment(ctx, &pb.CreateFragmentRequest{ - Name: name, - Config: cfg, - OrganizationId: orgID, - Visibility: &pbFragmentVisibility, - }) - if err != nil { - return nil, err - } - return fragmentFromProto(resp.Fragment), nil -} - -// UpdateFragment updates a fragment. -func (c *AppClient) UpdateFragment( - ctx context.Context, id, name string, config map[string]interface{}, public *bool, visibility *FragmentVisibility, -) (*Fragment, error) { - cfg, err := protoutils.StructToStructPb(config) - if err != nil { - return nil, err - } - pbVisibility := fragmentVisibilityToProto(*visibility) - resp, err := c.client.UpdateFragment(ctx, &pb.UpdateFragmentRequest{ - Id: id, - Name: name, - Config: cfg, - Public: public, - Visibility: &pbVisibility, - }) - if err != nil { - return nil, err - } - return fragmentFromProto(resp.Fragment), nil -} - -// DeleteFragment deletes a fragment. -func (c *AppClient) DeleteFragment(ctx context.Context, id string) error { - _, err := c.client.DeleteFragment(ctx, &pb.DeleteFragmentRequest{ - Id: id, - }) - return err -} - -// ListMachineFragments gets top level and nested fragments for a amchine, as well as any other fragments specified by IDs. Additional -// fragments are useful when needing to view fragments that will be provisionally added to the machine alongside existing fragments. -func (c *AppClient) ListMachineFragments(ctx context.Context, machineID string, additionalFragmentIDs []string) ([]*Fragment, error) { - resp, err := c.client.ListMachineFragments(ctx, &pb.ListMachineFragmentsRequest{ - MachineId: machineID, - AdditionalFragmentIds: additionalFragmentIDs, - }) - if err != nil { - return nil, err - } - var fragments []*Fragment - for _, fragment := range resp.Fragments { - fragments = append(fragments, fragmentFromProto(fragment)) - } - return fragments, nil -} - -// GetFragmentHistory gets the fragment's history and the next page token. -func (c *AppClient) GetFragmentHistory( - ctx context.Context, id string, pageToken *string, pageLimit *int64, -) ([]*FragmentHistoryEntry, string, error) { - resp, err := c.client.GetFragmentHistory(ctx, &pb.GetFragmentHistoryRequest{ - Id: id, - PageToken: pageToken, - PageLimit: pageLimit, - }) - if err != nil { - return nil, "", err - } - var history []*FragmentHistoryEntry - for _, entry := range resp.History { - history = append(history, fragmentHistoryEntryFromProto(entry)) - } - return history, resp.NextPageToken, nil -} - -// AddRole creates an identity authorization. -func (c *AppClient) AddRole(ctx context.Context, orgID, identityID, role, resourceType, resourceID string) error { - authorization, err := createAuthorization(orgID, identityID, "", role, resourceType, resourceID) - if err != nil { - return err - } - _, err = c.client.AddRole(ctx, &pb.AddRoleRequest{ - Authorization: authorization, - }) - return err -} - -// RemoveRole deletes an identity authorization. -func (c *AppClient) RemoveRole(ctx context.Context, orgID, identityID, role, resourceType, resourceID string) error { - authorization, err := createAuthorization(orgID, identityID, "", role, resourceType, resourceID) - if err != nil { - return err - } - _, err = c.client.RemoveRole(ctx, &pb.RemoveRoleRequest{ - Authorization: authorization, - }) - return err -} - -// ChangeRole changes an identity authorization to a new identity authorization. -func (c *AppClient) ChangeRole( - ctx context.Context, - oldOrgID, - oldIdentityID, - oldRole, - oldResourceType, - oldResourceID, - newOrgID, - newIdentityID, - newRole, - newResourceType, - newResourceID string, -) error { - oldAuthorization, err := createAuthorization(oldOrgID, oldIdentityID, "", oldRole, oldResourceType, oldResourceID) - if err != nil { - return err - } - newAuthorization, err := createAuthorization(newOrgID, newIdentityID, "", newRole, newResourceType, newResourceID) - if err != nil { - return err - } - _, err = c.client.ChangeRole(ctx, &pb.ChangeRoleRequest{ - OldAuthorization: oldAuthorization, - NewAuthorization: newAuthorization, - }) - return err -} - -// ListAuthorizations returns all authorization roles for any given resources. -// If no resources are given, all resources within the organization will be included. -func (c *AppClient) ListAuthorizations(ctx context.Context, orgID string, resourceIDs []string) ([]*Authorization, error) { - resp, err := c.client.ListAuthorizations(ctx, &pb.ListAuthorizationsRequest{ - OrganizationId: orgID, - ResourceIds: resourceIDs, - }) - if err != nil { - return nil, err - } - var authorizations []*Authorization - for _, authorization := range resp.Authorizations { - authorizations = append(authorizations, authorizationFromProto(authorization)) - } - return authorizations, nil -} - -// CheckPermissions checks the validity of a list of permissions. -func (c *AppClient) CheckPermissions(ctx context.Context, permissions []*AuthorizedPermissions) ([]*AuthorizedPermissions, error) { - var pbPermissions []*pb.AuthorizedPermissions - for _, permission := range permissions { - pbPermissions = append(pbPermissions, authorizedPermissionsToProto(permission)) - } - - resp, err := c.client.CheckPermissions(ctx, &pb.CheckPermissionsRequest{ - Permissions: pbPermissions, - }) - if err != nil { - return nil, err - } - - var authorizedPermissions []*AuthorizedPermissions - for _, permission := range resp.AuthorizedPermissions { - authorizedPermissions = append(authorizedPermissions, authorizedPermissionsFromProto(permission)) - } - return authorizedPermissions, nil -} - -// GetRegistryItem gets a registry item. -func (c *AppClient) GetRegistryItem(ctx context.Context, itemID string) (*RegistryItem, error) { - resp, err := c.client.GetRegistryItem(ctx, &pb.GetRegistryItemRequest{ - ItemId: itemID, - }) - if err != nil { - return nil, err - } - item, err := registryItemFromProto(resp.Item) - if err != nil { - return nil, err - } - return item, nil -} - -// CreateRegistryItem creates a registry item. -func (c *AppClient) CreateRegistryItem(ctx context.Context, orgID, name string, packageType PackageType) error { - _, err := c.client.CreateRegistryItem(ctx, &pb.CreateRegistryItemRequest{ - OrganizationId: orgID, - Name: name, - Type: packageTypeToProto(packageType), - }) - return err -} - -// UpdateRegistryItem updates a registry item. -func (c *AppClient) UpdateRegistryItem( - ctx context.Context, itemID string, packageType PackageType, description string, visibility Visibility, url *string, -) error { - _, err := c.client.UpdateRegistryItem(ctx, &pb.UpdateRegistryItemRequest{ - ItemId: itemID, - Type: packageTypeToProto(packageType), - Description: description, - Visibility: visibilityToProto(visibility), - Url: url, - }) - return err -} - -// ListRegistryItems lists the registry items in an organization. -func (c *AppClient) ListRegistryItems( - ctx context.Context, - orgID *string, - types []PackageType, - visibilities []Visibility, - platforms []string, - statuses []RegistryItemStatus, - searchTerm, - pageToken *string, - publicNamespaces []string, -) ([]*RegistryItem, error) { - var pbTypes []packages.PackageType - for _, packageType := range types { - pbTypes = append(pbTypes, packageTypeToProto(packageType)) - } - var pbVisibilities []pb.Visibility - for _, visibility := range visibilities { - pbVisibilities = append(pbVisibilities, visibilityToProto(visibility)) - } - var pbStatuses []pb.RegistryItemStatus - for _, status := range statuses { - pbStatuses = append(pbStatuses, registryItemStatusToProto(status)) - } - resp, err := c.client.ListRegistryItems(ctx, &pb.ListRegistryItemsRequest{ - OrganizationId: orgID, - Types: pbTypes, - Visibilities: pbVisibilities, - Platforms: platforms, - Statuses: pbStatuses, - SearchTerm: searchTerm, - PageToken: pageToken, - PublicNamespaces: publicNamespaces, - }) - if err != nil { - return nil, err - } - var items []*RegistryItem - for _, item := range resp.Items { - i, err := registryItemFromProto(item) - if err != nil { - return nil, err - } - items = append(items, i) - } - return items, nil -} - -// DeleteRegistryItem deletes a registry item given an ID that is formatted as `prefix:name“ -// where `prefix“ is the owner's organization ID or namespace. -func (c *AppClient) DeleteRegistryItem(ctx context.Context, itemID string) error { - _, err := c.client.DeleteRegistryItem(ctx, &pb.DeleteRegistryItemRequest{ - ItemId: itemID, - }) - return err -} - -// TransferRegistryItem transfers a registry item to a namespace. -func (c *AppClient) TransferRegistryItem(ctx context.Context, itemID, newPublicNamespace string) error { - _, err := c.client.TransferRegistryItem(ctx, &pb.TransferRegistryItemRequest{ - ItemId: itemID, - NewPublicNamespace: newPublicNamespace, - }) - return err -} - -// CreateModule creates a module and returns its ID and URL. -func (c *AppClient) CreateModule(ctx context.Context, orgID, name string) (string, string, error) { - resp, err := c.client.CreateModule(ctx, &pb.CreateModuleRequest{ - OrganizationId: orgID, - Name: name, - }) - if err != nil { - return "", "", err - } - return resp.ModuleId, resp.Url, nil -} - -// UpdateModule updates the documentation URL, description, models, entrypoint, and/or the visibility of a module and returns its URL. -// A path to a setup script can be added that is run before a newly downloaded module starts. -func (c *AppClient) UpdateModule( - ctx context.Context, moduleID string, visibility Visibility, url, description string, models []*Model, entrypoint string, firstRun *string, -) (string, error) { - var pbModels []*pb.Model - for _, model := range models { - pbModels = append(pbModels, modelToProto(model)) - } - resp, err := c.client.UpdateModule(ctx, &pb.UpdateModuleRequest{ - ModuleId: moduleID, - Visibility: visibilityToProto(visibility), - Url: url, - Description: description, - Models: pbModels, - Entrypoint: entrypoint, - FirstRun: firstRun, - }) - if err != nil { - return "", err - } - return resp.Url, nil -} - -// UploadModuleFile uploads a module file and returns the URL of the uploaded file. -func (c *AppClient) UploadModuleFile(ctx context.Context, fileInfo ModuleFileInfo, file []byte) (string, error) { - stream := &uploadModuleFileStream{client: c} - url, err := stream.startStream(ctx, &fileInfo, file) - if err != nil { - return "", err - } - - c.mu.Lock() - defer c.mu.Unlock() - return url, nil -} - -// GetModule gets a module. -func (c *AppClient) GetModule(ctx context.Context, moduleID string) (*Module, error) { - resp, err := c.client.GetModule(ctx, &pb.GetModuleRequest{ - ModuleId: moduleID, - }) - if err != nil { - return nil, err - } - return moduleFromProto(resp.Module), nil -} - -// ListModules lists the modules in the organization. -func (c *AppClient) ListModules(ctx context.Context, orgID *string) ([]*Module, error) { - resp, err := c.client.ListModules(ctx, &pb.ListModulesRequest{ - OrganizationId: orgID, - }) - if err != nil { - return nil, err - } - var modules []*Module - for _, module := range resp.Modules { - modules = append(modules, moduleFromProto(module)) - } - return modules, nil -} - -// CreateKey creates a new API key associated with a list of authorizations and returns its key and ID. -func (c *AppClient) CreateKey( - ctx context.Context, orgID string, keyAuthorizations []APIKeyAuthorization, name string, -) (string, string, error) { - var authorizations []*pb.Authorization - for _, keyAuthorization := range keyAuthorizations { - authorization, err := createAuthorization( - orgID, "", "api-key", keyAuthorization.role, keyAuthorization.resourceType, keyAuthorization.resourceID) - if err != nil { - return "", "", err - } - authorizations = append(authorizations, authorization) - } - - resp, err := c.client.CreateKey(ctx, &pb.CreateKeyRequest{ - Authorizations: authorizations, - Name: name, - }) - if err != nil { - return "", "", err - } - return resp.Key, resp.Id, nil -} - -// DeleteKey deletes an API key. -func (c *AppClient) DeleteKey(ctx context.Context, id string) error { - _, err := c.client.DeleteKey(ctx, &pb.DeleteKeyRequest{ - Id: id, - }) - return err -} - -// ListKeys lists all the keys for the organization. -func (c *AppClient) ListKeys(ctx context.Context, orgID string) ([]*APIKeyWithAuthorizations, error) { - resp, err := c.client.ListKeys(ctx, &pb.ListKeysRequest{ - OrgId: orgID, - }) - if err != nil { - return nil, err - } - var apiKeys []*APIKeyWithAuthorizations - for _, key := range resp.ApiKeys { - apiKeys = append(apiKeys, apiKeyWithAuthorizationsFromProto(key)) - } - return apiKeys, nil -} - -// RenameKey renames an API key and returns its ID and name. -func (c *AppClient) RenameKey(ctx context.Context, id, name string) (string, string, error) { - resp, err := c.client.RenameKey(ctx, &pb.RenameKeyRequest{ - Id: id, - Name: name, - }) - if err != nil { - return "", "", err - } - return resp.Id, resp.Name, nil -} - -// RotateKey rotates an API key and returns its ID and key. -func (c *AppClient) RotateKey(ctx context.Context, id string) (string, string, error) { - resp, err := c.client.RotateKey(ctx, &pb.RotateKeyRequest{ - Id: id, - }) - if err != nil { - return "", "", err - } - return resp.Id, resp.Key, nil -} - -// CreateKeyFromExistingKeyAuthorizations creates a new API key with an existing key's authorizations and returns its ID and key. -func (c *AppClient) CreateKeyFromExistingKeyAuthorizations(ctx context.Context, id string) (string, string, error) { - resp, err := c.client.CreateKeyFromExistingKeyAuthorizations(ctx, &pb.CreateKeyFromExistingKeyAuthorizationsRequest{ - Id: id, - }) - if err != nil { - return "", "", err - } - return resp.Id, resp.Key, nil -} From b1ce0d4e9a0cff8da295b547fe8be8a8412a8e1d Mon Sep 17 00:00:00 2001 From: purplenicole730 Date: Wed, 20 Nov 2024 14:47:40 -0500 Subject: [PATCH 05/15] add tests --- app/billing_client_test.go | 264 +++++++++++++++++++++ testutils/inject/billing_service_client.go | 74 ++++++ 2 files changed, 338 insertions(+) create mode 100644 app/billing_client_test.go create mode 100644 testutils/inject/billing_service_client.go diff --git a/app/billing_client_test.go b/app/billing_client_test.go new file mode 100644 index 00000000000..815d049fe8c --- /dev/null +++ b/app/billing_client_test.go @@ -0,0 +1,264 @@ +package app + +import ( + "context" + "errors" + "testing" + + pb "go.viam.com/api/app/v1" + "go.viam.com/rdk/testutils/inject" + "go.viam.com/test" + "google.golang.org/grpc" + "google.golang.org/protobuf/types/known/timestamppb" +) + +const ( + subtotal = 37 + sourceType = SourceTypeOrg + usageCostType = UsageCostTypeCloudStorage + cost float64 = 20 + discount float64 = 9 + totalWithDiscount = cost - discount + totalWithoutDiscount float64 = cost + email = "email" + paymentMethodType = PaymentMethodtypeCard + brand = "brand" + digits = "1234" + invoiceID = "invoice_id" + invoiceAmount float64 = 100.12 + status = "status" + balance float64 = 73.21 + billingOwnerOrgID = "billing_owner_organization_id" +) + +var ( + start = timestamppb.Timestamp{Seconds: 92, Nanos: 0} + end = timestamppb.Timestamp{Seconds: 99, Nanos: 999} + tier = "tier" + getCurrentMonthUsageResponse = GetCurrentMonthUsageResponse{ + StartDate: &start, + EndDate: &end, + ResourceUsageCostsBySource: []*ResourceUsageCostsBySource{ + { + SourceType: sourceType, + ResourceUsageCosts: &ResourceUsageCosts{ + UsageCosts: []*UsageCost{ + { + ResourceType: usageCostType, + Cost: cost, + }, + }, + Discount: discount, + TotalWithDiscount: totalWithDiscount, + TotalWithoutDiscount: totalWithoutDiscount, + }, + TierName: tier, + }, + }, + Subtotal: subtotal, + } + getOrgBillingInformationResponse = GetOrgBillingInformationResponse{ + Type: paymentMethodType, + BillingEmail: email, + Method: &PaymentMethodCard{ + Brand: brand, + LastFourDigits: digits, + }, + BillingTier: &tier, + } + invoiceDate = timestamppb.Timestamp{Seconds: 287, Nanos: 0} + dueDate = timestamppb.Timestamp{Seconds: 1241, Nanos: 40} + paidDate = timestamppb.Timestamp{Seconds: 827, Nanos: 62} + invoiceSummary = InvoiceSummary{ + ID: invoiceID, + InvoiceDate: &invoiceDate, + InvoiceAmount: invoiceAmount, + Status: status, + DueDate: &dueDate, + PaidDate: &paidDate, + } + chunk = []byte{4, 8} +) + +func sourceTypeToProto(sourceType SourceType) pb.SourceType { + switch sourceType { + case SourceTypeUnspecified: + return pb.SourceType_SOURCE_TYPE_UNSPECIFIED + case SourceTypeOrg: + return pb.SourceType_SOURCE_TYPE_ORG + case SourceTypeFragment: + return pb.SourceType_SOURCE_TYPE_FRAGMENT + default: + return pb.SourceType_SOURCE_TYPE_UNSPECIFIED + } +} + +func usageCostTypeToProto(costType UsageCostType) pb.UsageCostType { + switch costType { + case UsageCostTypeUnspecified: + return pb.UsageCostType_USAGE_COST_TYPE_UNSPECIFIED + case UsageCostTypeDataUpload: + return pb.UsageCostType_USAGE_COST_TYPE_DATA_UPLOAD + case UsageCostTypeDataEgress: + return pb.UsageCostType_USAGE_COST_TYPE_DATA_EGRESS + case UsageCostTypeRemoteControl: + return pb.UsageCostType_USAGE_COST_TYPE_REMOTE_CONTROL + case UsageCostTypeStandardCompute: + return pb.UsageCostType_USAGE_COST_TYPE_STANDARD_COMPUTE + case UsageCostTypeCloudStorage: + return pb.UsageCostType_USAGE_COST_TYPE_CLOUD_STORAGE + case UsageCostTypeBinaryDataCloudStorage: + return pb.UsageCostType_USAGE_COST_TYPE_BINARY_DATA_CLOUD_STORAGE + case UsageCostTypeOtherCloudStorage: + return pb.UsageCostType_USAGE_COST_TYPE_OTHER_CLOUD_STORAGE + case UsageCostTypePerMachine: + return pb.UsageCostType_USAGE_COST_TYPE_PER_MACHINE + default: + return pb.UsageCostType_USAGE_COST_TYPE_UNSPECIFIED + } +} + +func paymentMethodTypeToProto(methodType PaymentMethodType) pb.PaymentMethodType { + switch methodType { + case PaymentMethodTypeUnspecified: + return pb.PaymentMethodType_PAYMENT_METHOD_TYPE_UNSPECIFIED + case PaymentMethodtypeCard: + return pb.PaymentMethodType_PAYMENT_METHOD_TYPE_CARD + default: + return pb.PaymentMethodType_PAYMENT_METHOD_TYPE_UNSPECIFIED + } +} + +func createBillingGrpcClient() *inject.BillingServiceClient { + return &inject.BillingServiceClient{} +} + +type mockInvoiceStreamClient struct { + grpc.ClientStream + responses []*pb.GetInvoicePdfResponse + count int +} + +func (c *mockInvoiceStreamClient) Recv() (*pb.GetInvoicePdfResponse, error) { + if c.count >= len(c.responses) { + return nil, errors.New("end of reponses") + } + resp := c.responses[c.count] + c.count++ + return resp, nil +} + +func TestBillingClient(t *testing.T) { + grpcClient := createBillingGrpcClient() + client := BillingClient{client: grpcClient} + + t.Run("GetCurrentMonthUsage", func(t *testing.T) { + pbResponse := pb.GetCurrentMonthUsageResponse{ + StartDate: getCurrentMonthUsageResponse.StartDate, + EndDate: getCurrentMonthUsageResponse.EndDate, + ResourceUsageCostsBySource: []*pb.ResourceUsageCostsBySource{ + { + SourceType: sourceTypeToProto(sourceType), + ResourceUsageCosts: &pb.ResourceUsageCosts{ + UsageCosts: []*pb.UsageCost{ + { + ResourceType: usageCostTypeToProto(usageCostType), + Cost: cost, + }, + }, + Discount: discount, + TotalWithDiscount: totalWithDiscount, + TotalWithoutDiscount: totalWithoutDiscount, + }, + TierName: tier, + }, + }, + Subtotal: getCurrentMonthUsageResponse.Subtotal, + } + grpcClient.GetCurrentMonthUsageFunc = func( + ctx context.Context, in *pb.GetCurrentMonthUsageRequest, opts ...grpc.CallOption, + ) (*pb.GetCurrentMonthUsageResponse, error) { + test.That(t, in.OrgId, test.ShouldEqual, organizationID) + return &pbResponse, nil + } + resp, err := client.GetCurrentMonthUsage(context.Background(), organizationID) + test.That(t, err, test.ShouldBeNil) + test.That(t, resp, test.ShouldResemble, &getCurrentMonthUsageResponse) + }) + + t.Run("GetOrgBillingInformation", func(t *testing.T) { + pbResponse := pb.GetOrgBillingInformationResponse{ + Type: paymentMethodTypeToProto(getOrgBillingInformationResponse.Type), + BillingEmail: getOrgBillingInformationResponse.BillingEmail, + Method: &pb.PaymentMethodCard{ + Brand: getOrgBillingInformationResponse.Method.Brand, + LastFourDigits: getOrgBillingInformationResponse.Method.LastFourDigits, + }, + BillingTier: getOrgBillingInformationResponse.BillingTier, + } + grpcClient.GetOrgBillingInformationFunc = func( + ctx context.Context, in *pb.GetOrgBillingInformationRequest, opts ...grpc.CallOption, + ) (*pb.GetOrgBillingInformationResponse, error) { + test.That(t, in.OrgId, test.ShouldEqual, organizationID) + return &pbResponse, nil + } + resp, err := client.GetOrgBillingInformation(context.Background(), organizationID) + test.That(t, err, test.ShouldBeNil) + test.That(t, resp, test.ShouldResemble, &getOrgBillingInformationResponse) + }) + + t.Run("GetInvoicesSummary", func(t *testing.T) { + expectedInvoices := []*InvoiceSummary{&invoiceSummary} + grpcClient.GetInvoicesSummaryFunc = func( + ctx context.Context, in *pb.GetInvoicesSummaryRequest, opts ...grpc.CallOption, + ) (*pb.GetInvoicesSummaryResponse, error) { + test.That(t, in.OrgId, test.ShouldEqual, organizationID) + return &pb.GetInvoicesSummaryResponse{ + OutstandingBalance: balance, + Invoices: []*pb.InvoiceSummary{ + { + Id: invoiceSummary.ID, + InvoiceDate: invoiceSummary.InvoiceDate, + InvoiceAmount: invoiceSummary.InvoiceAmount, + Status: invoiceSummary.Status, + DueDate: invoiceSummary.DueDate, + PaidDate: invoiceSummary.PaidDate, + }, + }, + }, nil + } + outstandingBalance, invoices, err := client.GetInvoicesSummary(context.Background(), organizationID) + test.That(t, err, test.ShouldBeNil) + test.That(t, outstandingBalance, test.ShouldResemble, balance) + test.That(t, invoices, test.ShouldResemble, expectedInvoices) + }) + + t.Run("GetInvoicePDF", func(t *testing.T) { + ch := make(chan []byte) + grpcClient.GetInvoicePdfFunc = func( + ctx context.Context, in *pb.GetInvoicePdfRequest, opts ...grpc.CallOption, + ) (pb.BillingService_GetInvoicePdfClient, error) { + test.That(t, in.Id, test.ShouldEqual, invoiceID) + test.That(t, in.OrgId, test.ShouldEqual, organizationID) + return &mockInvoiceStreamClient{ + responses: []*pb.GetInvoicePdfResponse{ + {Chunk: chunk}, + }, + }, nil + } + err := client.GetInvoicePDF(context.Background(), invoiceID, organizationID, ch) + test.That(t, err, test.ShouldBeNil) + }) + + t.Run("SendPaymentRequiredEmail", func(t *testing.T) { + grpcClient.SendPaymentRequiredEmailFunc = func( + ctx context.Context, in *pb.SendPaymentRequiredEmailRequest, opts ...grpc.CallOption, + ) (*pb.SendPaymentRequiredEmailResponse, error) { + test.That(t, in.CustomerOrgId, test.ShouldEqual, organizationID) + test.That(t, in.BillingOwnerOrgId, test.ShouldEqual, billingOwnerOrgID) + return &pb.SendPaymentRequiredEmailResponse{}, nil + } + err := client.SendPaymentRequiredEmail(context.Background(), organizationID, billingOwnerOrgID) + test.That(t, err, test.ShouldBeNil) + }) +} diff --git a/testutils/inject/billing_service_client.go b/testutils/inject/billing_service_client.go new file mode 100644 index 00000000000..1bf2a918d53 --- /dev/null +++ b/testutils/inject/billing_service_client.go @@ -0,0 +1,74 @@ +package inject + +import ( + "context" + + pb "go.viam.com/api/app/v1" + "google.golang.org/grpc" +) + +// BillingServiceClient represents a fake instance of a billing service client. +type BillingServiceClient struct { + pb.BillingServiceClient + GetCurrentMonthUsageFunc func(ctx context.Context, in *pb.GetCurrentMonthUsageRequest, + opts ...grpc.CallOption) (*pb.GetCurrentMonthUsageResponse, error) + GetOrgBillingInformationFunc func(ctx context.Context, in *pb.GetOrgBillingInformationRequest, + opts ...grpc.CallOption) (*pb.GetOrgBillingInformationResponse, error) + GetInvoicesSummaryFunc func(ctx context.Context, in *pb.GetInvoicesSummaryRequest, + opts ...grpc.CallOption) (*pb.GetInvoicesSummaryResponse, error) + GetInvoicePdfFunc func(ctx context.Context, in *pb.GetInvoicePdfRequest, + opts ...grpc.CallOption) (pb.BillingService_GetInvoicePdfClient, error) + SendPaymentRequiredEmailFunc func(ctx context.Context, in *pb.SendPaymentRequiredEmailRequest, + opts ...grpc.CallOption) (*pb.SendPaymentRequiredEmailResponse, error) +} + + +// GetCurrentMonthUsage calls the injected GetCurrentMonthUsageFunc or the real verison. +func (bsc *BillingServiceClient) GetCurrentMonthUsage(ctx context.Context, in *pb.GetCurrentMonthUsageRequest, + opts ...grpc.CallOption, +) (*pb.GetCurrentMonthUsageResponse, error) { + if bsc.GetCurrentMonthUsageFunc == nil { + return bsc.BillingServiceClient.GetCurrentMonthUsage(ctx, in, opts...) + } + return bsc.GetCurrentMonthUsageFunc(ctx, in, opts...) +} + +// GetOrgBillingInformation calls the injected GetOrgBillingInformationFunc or the real verison. +func (bsc *BillingServiceClient) GetOrgBillingInformation(ctx context.Context, in *pb.GetOrgBillingInformationRequest, + opts ...grpc.CallOption, +) (*pb.GetOrgBillingInformationResponse, error) { + if bsc.GetOrgBillingInformationFunc == nil { + return bsc.BillingServiceClient.GetOrgBillingInformation(ctx, in, opts...) + } + return bsc.GetOrgBillingInformationFunc(ctx, in, opts...) +} + +// GetInvoicesSummary calls the injected GetInvoicesSummaryFunc or the real verison. +func (bsc *BillingServiceClient) GetInvoicesSummary(ctx context.Context, in *pb.GetInvoicesSummaryRequest, + opts ...grpc.CallOption, +) (*pb.GetInvoicesSummaryResponse, error) { + if bsc.GetInvoicesSummaryFunc == nil { + return bsc.BillingServiceClient.GetInvoicesSummary(ctx, in, opts...) + } + return bsc.GetInvoicesSummaryFunc(ctx, in, opts...) +} + +// GetInvoicePdf calls the injected GetInvoicePdfFunc or the real verison. +func (bsc *BillingServiceClient) GetInvoicePdf(ctx context.Context, in *pb.GetInvoicePdfRequest, + opts ...grpc.CallOption, +) (pb.BillingService_GetInvoicePdfClient, error) { + if bsc.GetInvoicePdfFunc == nil { + return bsc.BillingServiceClient.GetInvoicePdf(ctx, in, opts...) + } + return bsc.GetInvoicePdfFunc(ctx, in, opts...) +} + +// SendPaymentRequiredEmail calls the injected SendPaymentRequiredEmailFunc or the real verison. +func (bsc *BillingServiceClient) SendPaymentRequiredEmail(ctx context.Context, in *pb.SendPaymentRequiredEmailRequest, + opts ...grpc.CallOption, +) (*pb.SendPaymentRequiredEmailResponse, error) { + if bsc.SendPaymentRequiredEmailFunc == nil { + return bsc.BillingServiceClient.SendPaymentRequiredEmail(ctx, in, opts...) + } + return bsc.SendPaymentRequiredEmailFunc(ctx, in, opts...) +} From ac9be14b4610de86373df6a2b41a4386aeb25d7c Mon Sep 17 00:00:00 2001 From: purplenicole730 Date: Wed, 20 Nov 2024 15:02:37 -0500 Subject: [PATCH 06/15] add to viam client --- app/viam_client.go | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/app/viam_client.go b/app/viam_client.go index 6a8e3075012..4eb56ae7914 100644 --- a/app/viam_client.go +++ b/app/viam_client.go @@ -15,6 +15,8 @@ import ( // ViamClient is a gRPC client for method calls to Viam app. type ViamClient struct { conn rpc.ClientConn + logger logging.Logger + billingClient *BillingClient dataClient *DataClient } @@ -48,7 +50,7 @@ func CreateViamClientWithOptions(ctx context.Context, options Options, logger lo if err != nil { return nil, err } - return &ViamClient{conn: conn}, nil + return &ViamClient{conn: conn, logger: logger}, nil } // CreateViamClientWithAPIKey creates a ViamClient with an API key. @@ -63,6 +65,16 @@ func CreateViamClientWithAPIKey( return CreateViamClientWithOptions(ctx, options, logger) } +// Billingclient initializes and returns a Billingclient instance used to make app method calls. +// To use Billingclient, you must first instantiate a ViamClient. +func (c *ViamClient) Billingclient() *BillingClient { + if c.billingClient != nil { + return c.billingClient + } + c.billingClient = NewBillingClient(c.conn, c.logger) + return c.billingClient +} + // DataClient initializes and returns a DataClient instance used to make data method calls. // To use DataClient, you must first instantiate a ViamClient. func (c *ViamClient) DataClient() *DataClient { From 6e44722ae098678b76e11c1e94321d805b1d6f01 Mon Sep 17 00:00:00 2001 From: purplenicole730 Date: Wed, 20 Nov 2024 15:11:57 -0500 Subject: [PATCH 07/15] make lint --- app/billing_client.go | 93 +++++++++++----------- app/billing_client_test.go | 93 +++++++++++----------- app/viam_client.go | 6 +- testutils/inject/billing_service_client.go | 11 ++- 4 files changed, 101 insertions(+), 102 deletions(-) diff --git a/app/billing_client.go b/app/billing_client.go index fe5a2ef9c0a..07f446c1499 100644 --- a/app/billing_client.go +++ b/app/billing_client.go @@ -5,10 +5,11 @@ import ( "sync" pb "go.viam.com/api/app/v1" - "go.viam.com/rdk/logging" "go.viam.com/utils" "go.viam.com/utils/rpc" "google.golang.org/protobuf/types/known/timestamppb" + + "go.viam.com/rdk/logging" ) // UsageCostType specifies the type of usage cost. @@ -63,33 +64,33 @@ func usageCostTypeFromProto(costType pb.UsageCostType) UsageCostType { // UsageCost contains the cost and cost type. type UsageCost struct { ResourceType UsageCostType - Cost float64 + Cost float64 } func usageCostFromProto(cost *pb.UsageCost) *UsageCost { return &UsageCost{ ResourceType: usageCostTypeFromProto(cost.ResourceType), - Cost: cost.Cost, + Cost: cost.Cost, } } // ResourceUsageCosts holds the usage costs with discount information. type ResourceUsageCosts struct { - UsageCosts []*UsageCost - Discount float64 - TotalWithDiscount float64 + UsageCosts []*UsageCost + Discount float64 + TotalWithDiscount float64 TotalWithoutDiscount float64 } func resourceUsageCostsFromProto(costs *pb.ResourceUsageCosts) *ResourceUsageCosts { var usageCosts []*UsageCost - for _, cost := range(costs.UsageCosts) { + for _, cost := range costs.UsageCosts { usageCosts = append(usageCosts, usageCostFromProto(cost)) } return &ResourceUsageCosts{ - UsageCosts: usageCosts, - Discount: costs.Discount, - TotalWithDiscount: costs.TotalWithDiscount, + UsageCosts: usageCosts, + Discount: costs.Discount, + TotalWithDiscount: costs.TotalWithDiscount, TotalWithoutDiscount: costs.TotalWithoutDiscount, } } @@ -121,37 +122,37 @@ func sourceTypeFromProto(sourceType pb.SourceType) SourceType { // ResourceUsageCostsBySource contains the resource usage costs of a source type. type ResourceUsageCostsBySource struct { - SourceType SourceType + SourceType SourceType ResourceUsageCosts *ResourceUsageCosts - TierName string + TierName string } func resourceUsageCostsBySourceFromProto(costs *pb.ResourceUsageCostsBySource) *ResourceUsageCostsBySource { return &ResourceUsageCostsBySource{ - SourceType: sourceTypeFromProto(costs.SourceType), + SourceType: sourceTypeFromProto(costs.SourceType), ResourceUsageCosts: resourceUsageCostsFromProto(costs.ResourceUsageCosts), - TierName: costs.TierName, + TierName: costs.TierName, } } // GetCurrentMonthUsageResponse contains the current month usage information. type GetCurrentMonthUsageResponse struct { - StartDate *timestamppb.Timestamp - EndDate *timestamppb.Timestamp + StartDate *timestamppb.Timestamp + EndDate *timestamppb.Timestamp ResourceUsageCostsBySource []*ResourceUsageCostsBySource - Subtotal float64 + Subtotal float64 } func getCurrentMonthUsageResponseFromProto(response *pb.GetCurrentMonthUsageResponse) *GetCurrentMonthUsageResponse { var costs []*ResourceUsageCostsBySource - for _, cost := range(response.ResourceUsageCostsBySource) { + for _, cost := range response.ResourceUsageCostsBySource { costs = append(costs, resourceUsageCostsBySourceFromProto(cost)) } return &GetCurrentMonthUsageResponse{ - StartDate: response.StartDate, - EndDate: response.EndDate, + StartDate: response.StartDate, + EndDate: response.EndDate, ResourceUsageCostsBySource: costs, - Subtotal: response.Subtotal, + Subtotal: response.Subtotal, } } @@ -178,20 +179,20 @@ func paymentMethodTypeFromProto(methodType pb.PaymentMethodType) PaymentMethodTy // PaymentMethodCard holds the information of a card used for payment. type PaymentMethodCard struct { - Brand string + Brand string LastFourDigits string } func paymentMethodCardFromProto(card *pb.PaymentMethodCard) *PaymentMethodCard { return &PaymentMethodCard{ - Brand: card.Brand, + Brand: card.Brand, LastFourDigits: card.LastFourDigits, } } // GetOrgBillingInformationResponse contains the information of an organization's billing information. type GetOrgBillingInformationResponse struct { - Type PaymentMethodType + Type PaymentMethodType BillingEmail string // defined if type is PaymentMethodTypeCard Method *PaymentMethodCard @@ -201,38 +202,38 @@ type GetOrgBillingInformationResponse struct { func getOrgBillingInformationResponseFromProto(resp *pb.GetOrgBillingInformationResponse) *GetOrgBillingInformationResponse { return &GetOrgBillingInformationResponse{ - Type: paymentMethodTypeFromProto(resp.Type), + Type: paymentMethodTypeFromProto(resp.Type), BillingEmail: resp.BillingEmail, - Method: paymentMethodCardFromProto(resp.Method), - BillingTier: resp.BillingTier, + Method: paymentMethodCardFromProto(resp.Method), + BillingTier: resp.BillingTier, } } // InvoiceSummary holds the information of an invoice summary. type InvoiceSummary struct { - ID string - InvoiceDate *timestamppb.Timestamp + ID string + InvoiceDate *timestamppb.Timestamp InvoiceAmount float64 - Status string - DueDate *timestamppb.Timestamp - PaidDate *timestamppb.Timestamp + Status string + DueDate *timestamppb.Timestamp + PaidDate *timestamppb.Timestamp } func invoiceSummaryFromProto(summary *pb.InvoiceSummary) *InvoiceSummary { return &InvoiceSummary{ - ID: summary.Id, - InvoiceDate: summary.InvoiceDate, + ID: summary.Id, + InvoiceDate: summary.InvoiceDate, InvoiceAmount: summary.InvoiceAmount, - Status: summary.Status, - DueDate: summary.DueDate, - PaidDate: summary.PaidDate, + Status: summary.Status, + DueDate: summary.DueDate, + PaidDate: summary.PaidDate, } } type invoiceStream struct { - client *BillingClient + client *BillingClient streamCancel context.CancelFunc - streamMu sync.Mutex + streamMu sync.Mutex activeBackgroundWorkers sync.WaitGroup } @@ -249,7 +250,7 @@ func (s *invoiceStream) startStream(ctx context.Context, id, orgID string, ch ch s.streamCancel = cancel select { - case <- ctx.Done(): + case <-ctx.Done(): return ctx.Err() default: } @@ -257,7 +258,7 @@ func (s *invoiceStream) startStream(ctx context.Context, id, orgID string, ch ch // This call won't return any errors it had until the client tries to receive. //nolint:errcheck stream, _ := s.client.client.GetInvoicePdf(ctx, &pb.GetInvoicePdfRequest{ - Id: id, + Id: id, OrgId: orgID, }) _, err := stream.Recv() @@ -296,9 +297,7 @@ func (s *invoiceStream) receiveFromStream(ctx context.Context, stream pb.Billing } // If there is a response, send to the channel. var pdf []byte - for _, data := range streamResp.Chunk { - pdf = append(pdf, data) - } + pdf = append(pdf, streamResp.Chunk...) ch <- pdf } } @@ -347,14 +346,14 @@ func (c *BillingClient) GetInvoicesSummary(ctx context.Context, orgID string) (f return 0, nil, err } var invoices []*InvoiceSummary - for _, invoice := range(resp.Invoices) { + for _, invoice := range resp.Invoices { invoices = append(invoices, invoiceSummaryFromProto(invoice)) } return resp.OutstandingBalance, invoices, nil } // GetInvoicePDF gets the invoice PDF data. -func (c *BillingClient) GetInvoicePDF(ctx context.Context, id, orgID string, ch chan []byte) (error) { +func (c *BillingClient) GetInvoicePDF(ctx context.Context, id, orgID string, ch chan []byte) error { stream := &invoiceStream{client: c} err := stream.startStream(ctx, id, orgID, ch) @@ -370,7 +369,7 @@ func (c *BillingClient) GetInvoicePDF(ctx context.Context, id, orgID string, ch // SendPaymentRequiredEmail sends an email about payment requirement. func (c *BillingClient) SendPaymentRequiredEmail(ctx context.Context, customerOrgID, billingOwnerOrgID string) error { _, err := c.client.SendPaymentRequiredEmail(ctx, &pb.SendPaymentRequiredEmailRequest{ - CustomerOrgId: customerOrgID, + CustomerOrgId: customerOrgID, BillingOwnerOrgId: billingOwnerOrgID, }) return err diff --git a/app/billing_client_test.go b/app/billing_client_test.go index 815d049fe8c..597edfe2df7 100644 --- a/app/billing_client_test.go +++ b/app/billing_client_test.go @@ -6,38 +6,39 @@ import ( "testing" pb "go.viam.com/api/app/v1" - "go.viam.com/rdk/testutils/inject" "go.viam.com/test" "google.golang.org/grpc" "google.golang.org/protobuf/types/known/timestamppb" + + "go.viam.com/rdk/testutils/inject" ) const ( - subtotal = 37 - sourceType = SourceTypeOrg - usageCostType = UsageCostTypeCloudStorage - cost float64 = 20 - discount float64 = 9 - totalWithDiscount = cost - discount + subtotal = 37 + sourceType = SourceTypeOrg + usageCostType = UsageCostTypeCloudStorage + cost float64 = 20 + discount float64 = 9 + totalWithDiscount = cost - discount totalWithoutDiscount float64 = cost - email = "email" - paymentMethodType = PaymentMethodtypeCard - brand = "brand" - digits = "1234" - invoiceID = "invoice_id" - invoiceAmount float64 = 100.12 - status = "status" - balance float64 = 73.21 - billingOwnerOrgID = "billing_owner_organization_id" + email = "email" + paymentMethodType = PaymentMethodtypeCard + brand = "brand" + digits = "1234" + invoiceID = "invoice_id" + invoiceAmount float64 = 100.12 + status = "status" + balance float64 = 73.21 + billingOwnerOrgID = "billing_owner_organization_id" ) var ( - start = timestamppb.Timestamp{Seconds: 92, Nanos: 0} - end = timestamppb.Timestamp{Seconds: 99, Nanos: 999} - tier = "tier" + start = timestamppb.Timestamp{Seconds: 92, Nanos: 0} + end = timestamppb.Timestamp{Seconds: 99, Nanos: 999} + tier = "tier" getCurrentMonthUsageResponse = GetCurrentMonthUsageResponse{ StartDate: &start, - EndDate: &end, + EndDate: &end, ResourceUsageCostsBySource: []*ResourceUsageCostsBySource{ { SourceType: sourceType, @@ -45,11 +46,11 @@ var ( UsageCosts: []*UsageCost{ { ResourceType: usageCostType, - Cost: cost, + Cost: cost, }, }, - Discount: discount, - TotalWithDiscount: totalWithDiscount, + Discount: discount, + TotalWithDiscount: totalWithDiscount, TotalWithoutDiscount: totalWithoutDiscount, }, TierName: tier, @@ -58,24 +59,24 @@ var ( Subtotal: subtotal, } getOrgBillingInformationResponse = GetOrgBillingInformationResponse{ - Type: paymentMethodType, + Type: paymentMethodType, BillingEmail: email, Method: &PaymentMethodCard{ - Brand: brand, + Brand: brand, LastFourDigits: digits, }, BillingTier: &tier, } - invoiceDate = timestamppb.Timestamp{Seconds: 287, Nanos: 0} - dueDate = timestamppb.Timestamp{Seconds: 1241, Nanos: 40} - paidDate = timestamppb.Timestamp{Seconds: 827, Nanos: 62} + invoiceDate = timestamppb.Timestamp{Seconds: 287, Nanos: 0} + dueDate = timestamppb.Timestamp{Seconds: 1241, Nanos: 40} + paidDate = timestamppb.Timestamp{Seconds: 827, Nanos: 62} invoiceSummary = InvoiceSummary{ - ID: invoiceID, - InvoiceDate: &invoiceDate, + ID: invoiceID, + InvoiceDate: &invoiceDate, InvoiceAmount: invoiceAmount, - Status: status, - DueDate: &dueDate, - PaidDate: &paidDate, + Status: status, + DueDate: &dueDate, + PaidDate: &paidDate, } chunk = []byte{4, 8} ) @@ -136,7 +137,7 @@ func createBillingGrpcClient() *inject.BillingServiceClient { type mockInvoiceStreamClient struct { grpc.ClientStream responses []*pb.GetInvoicePdfResponse - count int + count int } func (c *mockInvoiceStreamClient) Recv() (*pb.GetInvoicePdfResponse, error) { @@ -155,7 +156,7 @@ func TestBillingClient(t *testing.T) { t.Run("GetCurrentMonthUsage", func(t *testing.T) { pbResponse := pb.GetCurrentMonthUsageResponse{ StartDate: getCurrentMonthUsageResponse.StartDate, - EndDate: getCurrentMonthUsageResponse.EndDate, + EndDate: getCurrentMonthUsageResponse.EndDate, ResourceUsageCostsBySource: []*pb.ResourceUsageCostsBySource{ { SourceType: sourceTypeToProto(sourceType), @@ -163,11 +164,11 @@ func TestBillingClient(t *testing.T) { UsageCosts: []*pb.UsageCost{ { ResourceType: usageCostTypeToProto(usageCostType), - Cost: cost, + Cost: cost, }, }, - Discount: discount, - TotalWithDiscount: totalWithDiscount, + Discount: discount, + TotalWithDiscount: totalWithDiscount, TotalWithoutDiscount: totalWithoutDiscount, }, TierName: tier, @@ -188,10 +189,10 @@ func TestBillingClient(t *testing.T) { t.Run("GetOrgBillingInformation", func(t *testing.T) { pbResponse := pb.GetOrgBillingInformationResponse{ - Type: paymentMethodTypeToProto(getOrgBillingInformationResponse.Type), + Type: paymentMethodTypeToProto(getOrgBillingInformationResponse.Type), BillingEmail: getOrgBillingInformationResponse.BillingEmail, Method: &pb.PaymentMethodCard{ - Brand: getOrgBillingInformationResponse.Method.Brand, + Brand: getOrgBillingInformationResponse.Method.Brand, LastFourDigits: getOrgBillingInformationResponse.Method.LastFourDigits, }, BillingTier: getOrgBillingInformationResponse.BillingTier, @@ -217,12 +218,12 @@ func TestBillingClient(t *testing.T) { OutstandingBalance: balance, Invoices: []*pb.InvoiceSummary{ { - Id: invoiceSummary.ID, - InvoiceDate: invoiceSummary.InvoiceDate, + Id: invoiceSummary.ID, + InvoiceDate: invoiceSummary.InvoiceDate, InvoiceAmount: invoiceSummary.InvoiceAmount, - Status: invoiceSummary.Status, - DueDate: invoiceSummary.DueDate, - PaidDate: invoiceSummary.PaidDate, + Status: invoiceSummary.Status, + DueDate: invoiceSummary.DueDate, + PaidDate: invoiceSummary.PaidDate, }, }, }, nil @@ -249,7 +250,7 @@ func TestBillingClient(t *testing.T) { err := client.GetInvoicePDF(context.Background(), invoiceID, organizationID, ch) test.That(t, err, test.ShouldBeNil) }) - + t.Run("SendPaymentRequiredEmail", func(t *testing.T) { grpcClient.SendPaymentRequiredEmailFunc = func( ctx context.Context, in *pb.SendPaymentRequiredEmailRequest, opts ...grpc.CallOption, diff --git a/app/viam_client.go b/app/viam_client.go index 4eb56ae7914..7f49e713e66 100644 --- a/app/viam_client.go +++ b/app/viam_client.go @@ -14,10 +14,10 @@ import ( // ViamClient is a gRPC client for method calls to Viam app. type ViamClient struct { - conn rpc.ClientConn - logger logging.Logger + conn rpc.ClientConn + logger logging.Logger billingClient *BillingClient - dataClient *DataClient + dataClient *DataClient } // Options has the options necessary to connect through gRPC. diff --git a/testutils/inject/billing_service_client.go b/testutils/inject/billing_service_client.go index 1bf2a918d53..a5cb97be47f 100644 --- a/testutils/inject/billing_service_client.go +++ b/testutils/inject/billing_service_client.go @@ -22,8 +22,7 @@ type BillingServiceClient struct { opts ...grpc.CallOption) (*pb.SendPaymentRequiredEmailResponse, error) } - -// GetCurrentMonthUsage calls the injected GetCurrentMonthUsageFunc or the real verison. +// GetCurrentMonthUsage calls the injected GetCurrentMonthUsageFunc or the real version. func (bsc *BillingServiceClient) GetCurrentMonthUsage(ctx context.Context, in *pb.GetCurrentMonthUsageRequest, opts ...grpc.CallOption, ) (*pb.GetCurrentMonthUsageResponse, error) { @@ -33,7 +32,7 @@ func (bsc *BillingServiceClient) GetCurrentMonthUsage(ctx context.Context, in *p return bsc.GetCurrentMonthUsageFunc(ctx, in, opts...) } -// GetOrgBillingInformation calls the injected GetOrgBillingInformationFunc or the real verison. +// GetOrgBillingInformation calls the injected GetOrgBillingInformationFunc or the real version. func (bsc *BillingServiceClient) GetOrgBillingInformation(ctx context.Context, in *pb.GetOrgBillingInformationRequest, opts ...grpc.CallOption, ) (*pb.GetOrgBillingInformationResponse, error) { @@ -43,7 +42,7 @@ func (bsc *BillingServiceClient) GetOrgBillingInformation(ctx context.Context, i return bsc.GetOrgBillingInformationFunc(ctx, in, opts...) } -// GetInvoicesSummary calls the injected GetInvoicesSummaryFunc or the real verison. +// GetInvoicesSummary calls the injected GetInvoicesSummaryFunc or the real version. func (bsc *BillingServiceClient) GetInvoicesSummary(ctx context.Context, in *pb.GetInvoicesSummaryRequest, opts ...grpc.CallOption, ) (*pb.GetInvoicesSummaryResponse, error) { @@ -53,7 +52,7 @@ func (bsc *BillingServiceClient) GetInvoicesSummary(ctx context.Context, in *pb. return bsc.GetInvoicesSummaryFunc(ctx, in, opts...) } -// GetInvoicePdf calls the injected GetInvoicePdfFunc or the real verison. +// GetInvoicePdf calls the injected GetInvoicePdfFunc or the real version. func (bsc *BillingServiceClient) GetInvoicePdf(ctx context.Context, in *pb.GetInvoicePdfRequest, opts ...grpc.CallOption, ) (pb.BillingService_GetInvoicePdfClient, error) { @@ -63,7 +62,7 @@ func (bsc *BillingServiceClient) GetInvoicePdf(ctx context.Context, in *pb.GetIn return bsc.GetInvoicePdfFunc(ctx, in, opts...) } -// SendPaymentRequiredEmail calls the injected SendPaymentRequiredEmailFunc or the real verison. +// SendPaymentRequiredEmail calls the injected SendPaymentRequiredEmailFunc or the real version. func (bsc *BillingServiceClient) SendPaymentRequiredEmail(ctx context.Context, in *pb.SendPaymentRequiredEmailRequest, opts ...grpc.CallOption, ) (*pb.SendPaymentRequiredEmailResponse, error) { From a65a213907f7e4e615a413bcd09db49933f9a4fa Mon Sep 17 00:00:00 2001 From: purplenicole730 Date: Thu, 21 Nov 2024 10:27:09 -0500 Subject: [PATCH 08/15] standardize inject file --- testutils/inject/billing_service_client.go | 44 +++++++++++----------- 1 file changed, 22 insertions(+), 22 deletions(-) diff --git a/testutils/inject/billing_service_client.go b/testutils/inject/billing_service_client.go index a5cb97be47f..b92b2699a80 100644 --- a/testutils/inject/billing_service_client.go +++ b/testutils/inject/billing_service_client.go @@ -3,29 +3,29 @@ package inject import ( "context" - pb "go.viam.com/api/app/v1" + billingpb "go.viam.com/api/app/v1" "google.golang.org/grpc" ) // BillingServiceClient represents a fake instance of a billing service client. type BillingServiceClient struct { - pb.BillingServiceClient - GetCurrentMonthUsageFunc func(ctx context.Context, in *pb.GetCurrentMonthUsageRequest, - opts ...grpc.CallOption) (*pb.GetCurrentMonthUsageResponse, error) - GetOrgBillingInformationFunc func(ctx context.Context, in *pb.GetOrgBillingInformationRequest, - opts ...grpc.CallOption) (*pb.GetOrgBillingInformationResponse, error) - GetInvoicesSummaryFunc func(ctx context.Context, in *pb.GetInvoicesSummaryRequest, - opts ...grpc.CallOption) (*pb.GetInvoicesSummaryResponse, error) - GetInvoicePdfFunc func(ctx context.Context, in *pb.GetInvoicePdfRequest, - opts ...grpc.CallOption) (pb.BillingService_GetInvoicePdfClient, error) - SendPaymentRequiredEmailFunc func(ctx context.Context, in *pb.SendPaymentRequiredEmailRequest, - opts ...grpc.CallOption) (*pb.SendPaymentRequiredEmailResponse, error) + billingpb.BillingServiceClient + GetCurrentMonthUsageFunc func(ctx context.Context, in *billingpb.GetCurrentMonthUsageRequest, + opts ...grpc.CallOption) (*billingpb.GetCurrentMonthUsageResponse, error) + GetOrgBillingInformationFunc func(ctx context.Context, in *billingpb.GetOrgBillingInformationRequest, + opts ...grpc.CallOption) (*billingpb.GetOrgBillingInformationResponse, error) + GetInvoicesSummaryFunc func(ctx context.Context, in *billingpb.GetInvoicesSummaryRequest, + opts ...grpc.CallOption) (*billingpb.GetInvoicesSummaryResponse, error) + GetInvoicePdfFunc func(ctx context.Context, in *billingpb.GetInvoicePdfRequest, + opts ...grpc.CallOption) (billingpb.BillingService_GetInvoicePdfClient, error) + SendPaymentRequiredEmailFunc func(ctx context.Context, in *billingpb.SendPaymentRequiredEmailRequest, + opts ...grpc.CallOption) (*billingpb.SendPaymentRequiredEmailResponse, error) } // GetCurrentMonthUsage calls the injected GetCurrentMonthUsageFunc or the real version. -func (bsc *BillingServiceClient) GetCurrentMonthUsage(ctx context.Context, in *pb.GetCurrentMonthUsageRequest, +func (bsc *BillingServiceClient) GetCurrentMonthUsage(ctx context.Context, in *billingpb.GetCurrentMonthUsageRequest, opts ...grpc.CallOption, -) (*pb.GetCurrentMonthUsageResponse, error) { +) (*billingpb.GetCurrentMonthUsageResponse, error) { if bsc.GetCurrentMonthUsageFunc == nil { return bsc.BillingServiceClient.GetCurrentMonthUsage(ctx, in, opts...) } @@ -33,9 +33,9 @@ func (bsc *BillingServiceClient) GetCurrentMonthUsage(ctx context.Context, in *p } // GetOrgBillingInformation calls the injected GetOrgBillingInformationFunc or the real version. -func (bsc *BillingServiceClient) GetOrgBillingInformation(ctx context.Context, in *pb.GetOrgBillingInformationRequest, +func (bsc *BillingServiceClient) GetOrgBillingInformation(ctx context.Context, in *billingpb.GetOrgBillingInformationRequest, opts ...grpc.CallOption, -) (*pb.GetOrgBillingInformationResponse, error) { +) (*billingpb.GetOrgBillingInformationResponse, error) { if bsc.GetOrgBillingInformationFunc == nil { return bsc.BillingServiceClient.GetOrgBillingInformation(ctx, in, opts...) } @@ -43,9 +43,9 @@ func (bsc *BillingServiceClient) GetOrgBillingInformation(ctx context.Context, i } // GetInvoicesSummary calls the injected GetInvoicesSummaryFunc or the real version. -func (bsc *BillingServiceClient) GetInvoicesSummary(ctx context.Context, in *pb.GetInvoicesSummaryRequest, +func (bsc *BillingServiceClient) GetInvoicesSummary(ctx context.Context, in *billingpb.GetInvoicesSummaryRequest, opts ...grpc.CallOption, -) (*pb.GetInvoicesSummaryResponse, error) { +) (*billingpb.GetInvoicesSummaryResponse, error) { if bsc.GetInvoicesSummaryFunc == nil { return bsc.BillingServiceClient.GetInvoicesSummary(ctx, in, opts...) } @@ -53,9 +53,9 @@ func (bsc *BillingServiceClient) GetInvoicesSummary(ctx context.Context, in *pb. } // GetInvoicePdf calls the injected GetInvoicePdfFunc or the real version. -func (bsc *BillingServiceClient) GetInvoicePdf(ctx context.Context, in *pb.GetInvoicePdfRequest, +func (bsc *BillingServiceClient) GetInvoicePdf(ctx context.Context, in *billingpb.GetInvoicePdfRequest, opts ...grpc.CallOption, -) (pb.BillingService_GetInvoicePdfClient, error) { +) (billingpb.BillingService_GetInvoicePdfClient, error) { if bsc.GetInvoicePdfFunc == nil { return bsc.BillingServiceClient.GetInvoicePdf(ctx, in, opts...) } @@ -63,9 +63,9 @@ func (bsc *BillingServiceClient) GetInvoicePdf(ctx context.Context, in *pb.GetIn } // SendPaymentRequiredEmail calls the injected SendPaymentRequiredEmailFunc or the real version. -func (bsc *BillingServiceClient) SendPaymentRequiredEmail(ctx context.Context, in *pb.SendPaymentRequiredEmailRequest, +func (bsc *BillingServiceClient) SendPaymentRequiredEmail(ctx context.Context, in *billingpb.SendPaymentRequiredEmailRequest, opts ...grpc.CallOption, -) (*pb.SendPaymentRequiredEmailResponse, error) { +) (*billingpb.SendPaymentRequiredEmailResponse, error) { if bsc.SendPaymentRequiredEmailFunc == nil { return bsc.BillingServiceClient.SendPaymentRequiredEmail(ctx, in, opts...) } From 20a73671aacff16776f6dda1036a7b1325f83bc5 Mon Sep 17 00:00:00 2001 From: purplenicole730 Date: Thu, 21 Nov 2024 13:31:04 -0500 Subject: [PATCH 09/15] change streaming test --- app/billing_client_test.go | 34 +++++++++------------- testutils/inject/billing_service_client.go | 13 +++++++++ 2 files changed, 26 insertions(+), 21 deletions(-) diff --git a/app/billing_client_test.go b/app/billing_client_test.go index 597edfe2df7..051fa49cabb 100644 --- a/app/billing_client_test.go +++ b/app/billing_client_test.go @@ -2,7 +2,6 @@ package app import ( "context" - "errors" "testing" pb "go.viam.com/api/app/v1" @@ -134,21 +133,6 @@ func createBillingGrpcClient() *inject.BillingServiceClient { return &inject.BillingServiceClient{} } -type mockInvoiceStreamClient struct { - grpc.ClientStream - responses []*pb.GetInvoicePdfResponse - count int -} - -func (c *mockInvoiceStreamClient) Recv() (*pb.GetInvoicePdfResponse, error) { - if c.count >= len(c.responses) { - return nil, errors.New("end of reponses") - } - resp := c.responses[c.count] - c.count++ - return resp, nil -} - func TestBillingClient(t *testing.T) { grpcClient := createBillingGrpcClient() client := BillingClient{client: grpcClient} @@ -235,20 +219,28 @@ func TestBillingClient(t *testing.T) { }) t.Run("GetInvoicePDF", func(t *testing.T) { + mockStream := &inject.BillingServiceGetInvoicePdfClient{ + RecvFunc: func() (*pb.GetInvoicePdfResponse, error) { + return &pb.GetInvoicePdfResponse{ + Chunk: chunk, + }, nil + }, + } ch := make(chan []byte) grpcClient.GetInvoicePdfFunc = func( ctx context.Context, in *pb.GetInvoicePdfRequest, opts ...grpc.CallOption, ) (pb.BillingService_GetInvoicePdfClient, error) { test.That(t, in.Id, test.ShouldEqual, invoiceID) test.That(t, in.OrgId, test.ShouldEqual, organizationID) - return &mockInvoiceStreamClient{ - responses: []*pb.GetInvoicePdfResponse{ - {Chunk: chunk}, - }, - }, nil + return mockStream, nil } err := client.GetInvoicePDF(context.Background(), invoiceID, organizationID, ch) test.That(t, err, test.ShouldBeNil) + // var resp []byte + // for chunkByte := range ch { + // resp = append(resp, chunkByte...) + // } + // test.That(t, resp, test.ShouldResemble, chunk) }) t.Run("SendPaymentRequiredEmail", func(t *testing.T) { diff --git a/testutils/inject/billing_service_client.go b/testutils/inject/billing_service_client.go index b92b2699a80..2d30f166e9d 100644 --- a/testutils/inject/billing_service_client.go +++ b/testutils/inject/billing_service_client.go @@ -62,6 +62,19 @@ func (bsc *BillingServiceClient) GetInvoicePdf(ctx context.Context, in *billingp return bsc.GetInvoicePdfFunc(ctx, in, opts...) } +// BillingServiceGetInvoicePdfClient represents a fake instance of a proto BillingService_GetInvoicePdfClient. +type BillingServiceGetInvoicePdfClient struct { + billingpb.BillingService_GetInvoicePdfClient + RecvFunc func() (*billingpb.GetInvoicePdfResponse, error) +} + +func (c *BillingServiceGetInvoicePdfClient) Recv() (*billingpb.GetInvoicePdfResponse, error) { + if c.RecvFunc == nil { + return c.BillingService_GetInvoicePdfClient.Recv() + } + return c.RecvFunc() +} + // SendPaymentRequiredEmail calls the injected SendPaymentRequiredEmailFunc or the real version. func (bsc *BillingServiceClient) SendPaymentRequiredEmail(ctx context.Context, in *billingpb.SendPaymentRequiredEmailRequest, opts ...grpc.CallOption, From 4073cc8e6d86eec6d1609dc493a4d99be26ca811 Mon Sep 17 00:00:00 2001 From: purplenicole730 Date: Thu, 21 Nov 2024 13:32:57 -0500 Subject: [PATCH 10/15] make lint --- testutils/inject/billing_service_client.go | 1 + 1 file changed, 1 insertion(+) diff --git a/testutils/inject/billing_service_client.go b/testutils/inject/billing_service_client.go index 2d30f166e9d..d7982b29749 100644 --- a/testutils/inject/billing_service_client.go +++ b/testutils/inject/billing_service_client.go @@ -68,6 +68,7 @@ type BillingServiceGetInvoicePdfClient struct { RecvFunc func() (*billingpb.GetInvoicePdfResponse, error) } +// Recv calls the injected RecvFunc or the real version. func (c *BillingServiceGetInvoicePdfClient) Recv() (*billingpb.GetInvoicePdfResponse, error) { if c.RecvFunc == nil { return c.BillingService_GetInvoicePdfClient.Recv() From d94b8a304e40cbb1dac997b93e7b31fab17a409b Mon Sep 17 00:00:00 2001 From: purplenicole730 Date: Thu, 21 Nov 2024 18:34:35 -0500 Subject: [PATCH 11/15] change getpdf --- app/billing_client.go | 112 ++++++++----------------------------- app/billing_client_test.go | 26 ++++++--- app/viam_client.go | 5 +- 3 files changed, 43 insertions(+), 100 deletions(-) diff --git a/app/billing_client.go b/app/billing_client.go index 07f446c1499..5022ed9f710 100644 --- a/app/billing_client.go +++ b/app/billing_client.go @@ -2,14 +2,12 @@ package app import ( "context" - "sync" + "errors" + "io" pb "go.viam.com/api/app/v1" - "go.viam.com/utils" "go.viam.com/utils/rpc" "google.golang.org/protobuf/types/known/timestamppb" - - "go.viam.com/rdk/logging" ) // UsageCostType specifies the type of usage cost. @@ -230,89 +228,14 @@ func invoiceSummaryFromProto(summary *pb.InvoiceSummary) *InvoiceSummary { } } -type invoiceStream struct { - client *BillingClient - streamCancel context.CancelFunc - streamMu sync.Mutex - - activeBackgroundWorkers sync.WaitGroup -} - -func (s *invoiceStream) startStream(ctx context.Context, id, orgID string, ch chan []byte) error { - s.streamMu.Lock() - defer s.streamMu.Unlock() - - if ctx.Err() != nil { - return ctx.Err() - } - - ctx, cancel := context.WithCancel(ctx) - s.streamCancel = cancel - - select { - case <-ctx.Done(): - return ctx.Err() - default: - } - - // This call won't return any errors it had until the client tries to receive. - //nolint:errcheck - stream, _ := s.client.client.GetInvoicePdf(ctx, &pb.GetInvoicePdfRequest{ - Id: id, - OrgId: orgID, - }) - _, err := stream.Recv() - if err != nil { - s.client.logger.CError(ctx, err) - return err - } - - // Create a background go routine to receive from the server stream. - // We rely on calling the Done function here rather than in close stream - // since managed go calls that function when the routine exits. - s.activeBackgroundWorkers.Add(1) - utils.ManagedGo(func() { - s.receiveFromStream(ctx, stream, ch) - }, - s.activeBackgroundWorkers.Done) - return nil -} - -func (s *invoiceStream) receiveFromStream(ctx context.Context, stream pb.BillingService_GetInvoicePdfClient, ch chan []byte) { - defer s.streamCancel() - - // repeatedly receive from the stream - for { - select { - case <-ctx.Done(): - s.client.logger.Debug(ctx.Err()) - return - default: - } - streamResp, err := stream.Recv() - if err != nil { - // only debug log the context canceled error - s.client.logger.Debug(err) - return - } - // If there is a response, send to the channel. - var pdf []byte - pdf = append(pdf, streamResp.Chunk...) - ch <- pdf - } -} - // BillingClient is a gRPC client for method calls to the Billing API. type BillingClient struct { client pb.BillingServiceClient - logger logging.Logger - - mu sync.Mutex } // NewBillingClient constructs a new BillingClient using the connection passed in by the Viam client. -func NewBillingClient(conn rpc.ClientConn, logger logging.Logger) *BillingClient { - return &BillingClient{client: pb.NewBillingServiceClient(conn), logger: logger} +func NewBillingClient(conn rpc.ClientConn) *BillingClient { + return &BillingClient{client: pb.NewBillingServiceClient(conn)} } // GetCurrentMonthUsage gets the data usage information for the current month for an organization. @@ -353,17 +276,28 @@ func (c *BillingClient) GetInvoicesSummary(ctx context.Context, orgID string) (f } // GetInvoicePDF gets the invoice PDF data. -func (c *BillingClient) GetInvoicePDF(ctx context.Context, id, orgID string, ch chan []byte) error { - stream := &invoiceStream{client: c} - - err := stream.startStream(ctx, id, orgID, ch) +func (c *BillingClient) GetInvoicePDF(ctx context.Context, id, orgID string) ([]byte, error) { + stream, err := c.client.GetInvoicePdf(ctx, &pb.GetInvoicePdfRequest{ + Id: id, + OrgId: orgID, + }) if err != nil { - return nil + return nil, err + } + + var data []byte + for { + resp, err := stream.Recv() + if err != nil { + if errors.Is(err, io.EOF) { + break + } + return data, err + } + data = append(data, resp.Chunk...) } - c.mu.Lock() - defer c.mu.Unlock() - return nil + return data, nil } // SendPaymentRequiredEmail sends an email about payment requirement. diff --git a/app/billing_client_test.go b/app/billing_client_test.go index 051fa49cabb..ac2bb83de15 100644 --- a/app/billing_client_test.go +++ b/app/billing_client_test.go @@ -2,6 +2,7 @@ package app import ( "context" + "io" "testing" pb "go.viam.com/api/app/v1" @@ -77,7 +78,11 @@ var ( DueDate: &dueDate, PaidDate: &paidDate, } - chunk = []byte{4, 8} + chunk1 = []byte{4, 8} + chunk2 = []byte("chunk1") + chunk3 = []byte("chunk2") + chunks = [][]byte{chunk1, chunk2, chunk3} + chunkCount = len(chunks) ) func sourceTypeToProto(sourceType SourceType) pb.SourceType { @@ -219,14 +224,23 @@ func TestBillingClient(t *testing.T) { }) t.Run("GetInvoicePDF", func(t *testing.T) { + var expectedData []byte + expectedData = append(expectedData, chunk1...) + expectedData = append(expectedData, chunk2...) + expectedData = append(expectedData, chunk3...) + var count int mockStream := &inject.BillingServiceGetInvoicePdfClient{ RecvFunc: func() (*pb.GetInvoicePdfResponse, error) { + if count >= chunkCount { + return nil, io.EOF + } + chunk := chunks[count] + count++ return &pb.GetInvoicePdfResponse{ Chunk: chunk, }, nil }, } - ch := make(chan []byte) grpcClient.GetInvoicePdfFunc = func( ctx context.Context, in *pb.GetInvoicePdfRequest, opts ...grpc.CallOption, ) (pb.BillingService_GetInvoicePdfClient, error) { @@ -234,13 +248,9 @@ func TestBillingClient(t *testing.T) { test.That(t, in.OrgId, test.ShouldEqual, organizationID) return mockStream, nil } - err := client.GetInvoicePDF(context.Background(), invoiceID, organizationID, ch) + data, err := client.GetInvoicePDF(context.Background(), invoiceID, organizationID) test.That(t, err, test.ShouldBeNil) - // var resp []byte - // for chunkByte := range ch { - // resp = append(resp, chunkByte...) - // } - // test.That(t, resp, test.ShouldResemble, chunk) + test.That(t, data, test.ShouldResemble, expectedData) }) t.Run("SendPaymentRequiredEmail", func(t *testing.T) { diff --git a/app/viam_client.go b/app/viam_client.go index 7f49e713e66..74e155ed46a 100644 --- a/app/viam_client.go +++ b/app/viam_client.go @@ -15,7 +15,6 @@ import ( // ViamClient is a gRPC client for method calls to Viam app. type ViamClient struct { conn rpc.ClientConn - logger logging.Logger billingClient *BillingClient dataClient *DataClient } @@ -50,7 +49,7 @@ func CreateViamClientWithOptions(ctx context.Context, options Options, logger lo if err != nil { return nil, err } - return &ViamClient{conn: conn, logger: logger}, nil + return &ViamClient{conn: conn}, nil } // CreateViamClientWithAPIKey creates a ViamClient with an API key. @@ -71,7 +70,7 @@ func (c *ViamClient) Billingclient() *BillingClient { if c.billingClient != nil { return c.billingClient } - c.billingClient = NewBillingClient(c.conn, c.logger) + c.billingClient = NewBillingClient(c.conn) return c.billingClient } From b8482dc19e9965505aabae9fb609bc76f352a586 Mon Sep 17 00:00:00 2001 From: purplenicole730 Date: Thu, 21 Nov 2024 18:38:24 -0500 Subject: [PATCH 12/15] add billing client to viam client test --- app/viam_client_test.go | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/app/viam_client_test.go b/app/viam_client_test.go index 08b93ad1c31..908b7c141ab 100644 --- a/app/viam_client_test.go +++ b/app/viam_client_test.go @@ -119,7 +119,7 @@ func TestCreateViamClientWithAPIKeyTests(t *testing.T) { } } -func TestNewDataClient(t *testing.T) { +func TestNewAppClients(t *testing.T) { originalDialDirectGRPC := dialDirectGRPC dialDirectGRPC = mockDialDirectGRPC defer func() { dialDirectGRPC = originalDialDirectGRPC }() @@ -135,6 +135,16 @@ func TestNewDataClient(t *testing.T) { test.That(t, err, test.ShouldBeNil) defer client.Close() + billingClient := client.Billingclient() + test.That(t, billingClient, test.ShouldNotBeNil) + test.That(t, billingClient, test.ShouldHaveSameTypeAs, &BillingClient{}) + test.That(t, billingClient.client, test.ShouldImplement, (*pb.DataServiceClient)(nil)) + + // Testing that a second call to Billingclient() returns the same instance + billingClient2 := client.Billingclient() + test.That(t, billingClient2, test.ShouldNotBeNil) + test.That(t, billingClient, test.ShouldResemble, billingClient2) + dataClient := client.DataClient() test.That(t, dataClient, test.ShouldNotBeNil) test.That(t, dataClient, test.ShouldHaveSameTypeAs, &DataClient{}) From 44400345224b7130966e650075479258a81208aa Mon Sep 17 00:00:00 2001 From: purplenicole730 Date: Thu, 21 Nov 2024 18:47:01 -0500 Subject: [PATCH 13/15] add billing test to viam client test --- app/viam_client_test.go | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/app/viam_client_test.go b/app/viam_client_test.go index 908b7c141ab..6bd4eaa574f 100644 --- a/app/viam_client_test.go +++ b/app/viam_client_test.go @@ -5,7 +5,8 @@ import ( "testing" "github.com/viamrobotics/webrtc/v3" - pb "go.viam.com/api/app/data/v1" + datapb "go.viam.com/api/app/data/v1" + apppb "go.viam.com/api/app/v1" "go.viam.com/test" "go.viam.com/utils" "go.viam.com/utils/rpc" @@ -138,17 +139,17 @@ func TestNewAppClients(t *testing.T) { billingClient := client.Billingclient() test.That(t, billingClient, test.ShouldNotBeNil) test.That(t, billingClient, test.ShouldHaveSameTypeAs, &BillingClient{}) - test.That(t, billingClient.client, test.ShouldImplement, (*pb.DataServiceClient)(nil)) + test.That(t, billingClient.client, test.ShouldImplement, (*apppb.BillingServiceClient)(nil)) // Testing that a second call to Billingclient() returns the same instance billingClient2 := client.Billingclient() test.That(t, billingClient2, test.ShouldNotBeNil) - test.That(t, billingClient, test.ShouldResemble, billingClient2) + test.That(t, billingClient, test.ShouldEqual, billingClient2) dataClient := client.DataClient() test.That(t, dataClient, test.ShouldNotBeNil) test.That(t, dataClient, test.ShouldHaveSameTypeAs, &DataClient{}) - test.That(t, dataClient.client, test.ShouldImplement, (*pb.DataServiceClient)(nil)) + test.That(t, dataClient.client, test.ShouldImplement, (*datapb.DataServiceClient)(nil)) // Testing that a second call to DataClient() returns the same instance dataClient2 := client.DataClient() From a09687cfb43e658d00c5073f9cfc25ecf92c5609 Mon Sep 17 00:00:00 2001 From: purplenicole730 Date: Thu, 21 Nov 2024 18:53:19 -0500 Subject: [PATCH 14/15] use bytes join --- app/billing_client_test.go | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/app/billing_client_test.go b/app/billing_client_test.go index ac2bb83de15..e8af00302a0 100644 --- a/app/billing_client_test.go +++ b/app/billing_client_test.go @@ -1,6 +1,7 @@ package app import ( + "bytes" "context" "io" "testing" @@ -224,10 +225,7 @@ func TestBillingClient(t *testing.T) { }) t.Run("GetInvoicePDF", func(t *testing.T) { - var expectedData []byte - expectedData = append(expectedData, chunk1...) - expectedData = append(expectedData, chunk2...) - expectedData = append(expectedData, chunk3...) + expectedData := bytes.Join(chunks, nil) var count int mockStream := &inject.BillingServiceGetInvoicePdfClient{ RecvFunc: func() (*pb.GetInvoicePdfResponse, error) { From abaabda643cbb6790f5628473f89e931528ce88d Mon Sep 17 00:00:00 2001 From: purplenicole730 Date: Thu, 21 Nov 2024 18:56:27 -0500 Subject: [PATCH 15/15] fix comment --- app/billing_client.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/app/billing_client.go b/app/billing_client.go index 5022ed9f710..299c4f6fd35 100644 --- a/app/billing_client.go +++ b/app/billing_client.go @@ -118,7 +118,7 @@ func sourceTypeFromProto(sourceType pb.SourceType) SourceType { } } -// ResourceUsageCostsBySource contains the resource usage costs of a source type. +// ResourceUsageCostsBySource contains the resource usage costs of a source. type ResourceUsageCostsBySource struct { SourceType SourceType ResourceUsageCosts *ResourceUsageCosts