Add AWS Metering service support (#15290)
Этот коммит содержится в:
catalintomai
коммит произвёл
GitHub
GitHub
родитель
50e37068b5
Коммит
1c0d590c81
155
services/awsmeter/awsmeter.go
Обычный файл
155
services/awsmeter/awsmeter.go
Обычный файл
@@ -0,0 +1,155 @@
|
||||
// Copyright (c) 2015-present Mattermost, Inc. All Rights Reserved.
|
||||
// See LICENSE.txt for license information.
|
||||
|
||||
package awsmeter
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"os"
|
||||
"time"
|
||||
|
||||
"github.com/mattermost/mattermost-server/v5/mlog"
|
||||
"github.com/mattermost/mattermost-server/v5/model"
|
||||
"github.com/mattermost/mattermost-server/v5/store"
|
||||
"github.com/pkg/errors"
|
||||
|
||||
"github.com/aws/aws-sdk-go/aws"
|
||||
"github.com/aws/aws-sdk-go/aws/credentials"
|
||||
"github.com/aws/aws-sdk-go/aws/credentials/ec2rolecreds"
|
||||
"github.com/aws/aws-sdk-go/aws/ec2metadata"
|
||||
"github.com/aws/aws-sdk-go/aws/session"
|
||||
"github.com/aws/aws-sdk-go/service/marketplacemetering"
|
||||
"github.com/aws/aws-sdk-go/service/marketplacemetering/marketplacemeteringiface"
|
||||
)
|
||||
|
||||
type AwsMeter struct {
|
||||
store store.Store
|
||||
service *AWSMeterService
|
||||
config *model.Config
|
||||
}
|
||||
|
||||
type AWSMeterService struct {
|
||||
AwsDryRun bool
|
||||
AwsProductCode string
|
||||
AwsMeteringSvc marketplacemeteringiface.MarketplaceMeteringAPI
|
||||
}
|
||||
|
||||
type AWSMeterReport struct {
|
||||
Dimension string `json:"dimension"`
|
||||
Value int64 `json:"value"`
|
||||
Timestamp time.Time `json:"timestamp"`
|
||||
}
|
||||
|
||||
func (o *AWSMeterReport) ToJSON() string {
|
||||
b, _ := json.Marshal(o)
|
||||
return string(b)
|
||||
}
|
||||
|
||||
func New(store store.Store, config *model.Config) *AwsMeter {
|
||||
svc := &AWSMeterService{
|
||||
AwsDryRun: false,
|
||||
AwsProductCode: "12345", //TODO
|
||||
}
|
||||
|
||||
service, err := newAWSMarketplaceMeteringService()
|
||||
if err != nil {
|
||||
mlog.Error("newAWSMeterService", mlog.String("error", err.Error()))
|
||||
return nil
|
||||
}
|
||||
|
||||
svc.AwsMeteringSvc = service
|
||||
return &AwsMeter{
|
||||
store: store,
|
||||
service: svc,
|
||||
config: config,
|
||||
}
|
||||
}
|
||||
|
||||
func newAWSMarketplaceMeteringService() (*marketplacemetering.MarketplaceMetering, error) {
|
||||
region := os.Getenv("AWS_REGION")
|
||||
s, err := session.NewSession(&aws.Config{Region: ®ion})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
creds := credentials.NewChainCredentials(
|
||||
[]credentials.Provider{
|
||||
&ec2rolecreds.EC2RoleProvider{
|
||||
Client: ec2metadata.New(s),
|
||||
},
|
||||
})
|
||||
|
||||
_, err = creds.Get()
|
||||
if err != nil {
|
||||
mlog.Error("session is invalid", mlog.String("error", err.Error()))
|
||||
return nil, errors.New("cannot obtain credentials")
|
||||
}
|
||||
|
||||
return marketplacemetering.New(session.Must(session.NewSession(&aws.Config{
|
||||
Credentials: creds,
|
||||
}))), nil
|
||||
}
|
||||
|
||||
// a report entry is for all metrics
|
||||
func (awsm *AwsMeter) GetUserCategoryUsage(dimensions []string, startTime time.Time, endTime time.Time) []*AWSMeterReport {
|
||||
reports := make([]*AWSMeterReport, 0)
|
||||
|
||||
for _, dimension := range dimensions {
|
||||
var userCount int64
|
||||
var err error
|
||||
|
||||
switch dimension {
|
||||
case model.AWS_METERING_DIMENSION_USAGE_HRS:
|
||||
userCount, err = awsm.store.User().AnalyticsActiveCountForPeriod(model.GetMillisForTime(startTime), model.GetMillisForTime(endTime), model.UserCountOptions{})
|
||||
if err != nil {
|
||||
mlog.Error("Failed to obtain usage data", mlog.String("dimension", dimension), mlog.String("start", startTime.String()), mlog.Int64("count", userCount), mlog.Err(err))
|
||||
continue
|
||||
}
|
||||
default:
|
||||
mlog.Error("Dimension does not exist!", mlog.String("dimension", dimension))
|
||||
continue
|
||||
}
|
||||
|
||||
report := &AWSMeterReport{
|
||||
Dimension: dimension,
|
||||
Value: userCount,
|
||||
Timestamp: startTime,
|
||||
}
|
||||
|
||||
reports = append(reports, report)
|
||||
}
|
||||
|
||||
return reports
|
||||
}
|
||||
|
||||
func (awsm *AwsMeter) ReportUserCategoryUsage(reports []*AWSMeterReport) error {
|
||||
for _, report := range reports {
|
||||
err := sendReportToMeteringService(awsm.service, report)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func sendReportToMeteringService(ams *AWSMeterService, report *AWSMeterReport) error {
|
||||
params := &marketplacemetering.MeterUsageInput{
|
||||
DryRun: aws.Bool(ams.AwsDryRun),
|
||||
ProductCode: aws.String(ams.AwsProductCode),
|
||||
UsageDimension: aws.String(report.Dimension),
|
||||
UsageQuantity: aws.Int64(report.Value),
|
||||
Timestamp: aws.Time(report.Timestamp),
|
||||
}
|
||||
|
||||
resp, err := ams.AwsMeteringSvc.MeterUsage(params)
|
||||
if err != nil {
|
||||
return errors.Wrap(err, "Invalid metering service id.")
|
||||
}
|
||||
if resp.MeteringRecordId == nil {
|
||||
return errors.Wrap(err, "Invalid metering service id.")
|
||||
}
|
||||
|
||||
mlog.Debug("Sent record to AWS metering service", mlog.String("dimension", report.Dimension), mlog.Int64("value", report.Value), mlog.String("timestamp", report.Timestamp.String()))
|
||||
|
||||
return nil
|
||||
}
|
||||
149
services/awsmeter/awsmeter_test.go
Обычный файл
149
services/awsmeter/awsmeter_test.go
Обычный файл
@@ -0,0 +1,149 @@
|
||||
// Copyright (c) 2015-present Mattermost, Inc. All Rights Reserved.
|
||||
// See LICENSE.txt for license information.
|
||||
|
||||
package awsmeter
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/mattermost/mattermost-server/v5/model"
|
||||
"github.com/mattermost/mattermost-server/v5/plugin/plugintest/mock"
|
||||
"github.com/mattermost/mattermost-server/v5/store/storetest/mocks"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/aws/aws-sdk-go/service/marketplacemetering"
|
||||
"github.com/aws/aws-sdk-go/service/marketplacemetering/marketplacemeteringiface"
|
||||
)
|
||||
|
||||
type mockMarketplaceMeteringClient struct {
|
||||
marketplacemeteringiface.MarketplaceMeteringAPI
|
||||
}
|
||||
|
||||
func (m *mockMarketplaceMeteringClient) MeterUsage(input *marketplacemetering.MeterUsageInput) (*marketplacemetering.MeterUsageOutput, error) {
|
||||
return &marketplacemetering.MeterUsageOutput{
|
||||
MeteringRecordId: String("1"),
|
||||
}, nil
|
||||
}
|
||||
|
||||
type mockMarketplaceMeteringClientWithError struct {
|
||||
marketplacemeteringiface.MarketplaceMeteringAPI
|
||||
}
|
||||
|
||||
func (m *mockMarketplaceMeteringClientWithError) MeterUsage(input *marketplacemetering.MeterUsageInput) (*marketplacemetering.MeterUsageOutput, error) {
|
||||
return nil, errors.New("error")
|
||||
}
|
||||
|
||||
func String(i string) *string {
|
||||
return &i
|
||||
}
|
||||
func TestAwsMeterUsage(t *testing.T) {
|
||||
startTime := time.Now()
|
||||
endTime := time.Now()
|
||||
dimensions := []string{model.AWS_METERING_DIMENSION_USAGE_HRS}
|
||||
|
||||
userStoreMock := mocks.UserStore{}
|
||||
userStoreMock.On("AnalyticsActiveCountForPeriod", model.GetMillisForTime(startTime), model.GetMillisForTime(endTime), mock.AnythingOfType("model.UserCountOptions")).Return(int64(2), nil)
|
||||
|
||||
storeMock := mocks.Store{}
|
||||
storeMock.On("User").Return(&userStoreMock)
|
||||
|
||||
reports := make([]*AWSMeterReport, 1)
|
||||
reports[0] = &AWSMeterReport{
|
||||
Dimension: model.AWS_METERING_DIMENSION_USAGE_HRS,
|
||||
Value: 2,
|
||||
Timestamp: startTime,
|
||||
}
|
||||
|
||||
// Define a mock struct to be used in your unit tests of myFunc.
|
||||
svc := &AWSMeterService{
|
||||
AwsDryRun: false,
|
||||
AwsProductCode: "12345",
|
||||
AwsMeteringSvc: &mockMarketplaceMeteringClient{},
|
||||
}
|
||||
|
||||
config := &model.Config{}
|
||||
config.SetDefaults()
|
||||
|
||||
awsmeter := &AwsMeter{
|
||||
store: &storeMock,
|
||||
service: svc,
|
||||
config: config,
|
||||
}
|
||||
|
||||
t.Run("Send report for one usage category", func(t *testing.T) {
|
||||
resultReports := awsmeter.GetUserCategoryUsage(dimensions, startTime, endTime)
|
||||
require.NotNil(t, resultReports)
|
||||
assert.Equal(t, 1, len(resultReports))
|
||||
assert.Equal(t, reports[0].Dimension, resultReports[0].Dimension)
|
||||
assert.Equal(t, reports[0].Value, resultReports[0].Value)
|
||||
assert.Equal(t, reports[0].Timestamp, resultReports[0].Timestamp)
|
||||
|
||||
err := awsmeter.ReportUserCategoryUsage(resultReports)
|
||||
require.Nil(t, err)
|
||||
})
|
||||
|
||||
t.Run("Error in AWS service call", func(t *testing.T) {
|
||||
awsmeter.service.AwsMeteringSvc = &mockMarketplaceMeteringClientWithError{}
|
||||
resultReports := awsmeter.GetUserCategoryUsage(dimensions, startTime, endTime)
|
||||
require.NotNil(t, resultReports)
|
||||
assert.Equal(t, 1, len(resultReports))
|
||||
err := awsmeter.ReportUserCategoryUsage(resultReports)
|
||||
require.NotNil(t, err)
|
||||
})
|
||||
|
||||
t.Run("Invalid dimension", func(t *testing.T) {
|
||||
awsmeter.service.AwsMeteringSvc = &mockMarketplaceMeteringClient{}
|
||||
dimensions = []string{"invalid dimension"}
|
||||
resultReports := awsmeter.GetUserCategoryUsage(dimensions, startTime, endTime)
|
||||
require.NotNil(t, resultReports)
|
||||
assert.Equal(t, 0, len(resultReports))
|
||||
err := awsmeter.ReportUserCategoryUsage(resultReports)
|
||||
require.Nil(t, err)
|
||||
})
|
||||
}
|
||||
|
||||
func TestAwsMeterUsageWithDBError(t *testing.T) {
|
||||
startTime := time.Now()
|
||||
endTime := time.Now()
|
||||
dimensions := []string{model.AWS_METERING_DIMENSION_USAGE_HRS}
|
||||
|
||||
userStoreMock := mocks.UserStore{}
|
||||
userStoreMock.On("AnalyticsActiveCountForPeriod", model.GetMillisForTime(startTime), model.GetMillisForTime(endTime), mock.AnythingOfType("model.UserCountOptions")).Return(int64(0), errors.New("error"))
|
||||
|
||||
storeMock := mocks.Store{}
|
||||
storeMock.On("User").Return(&userStoreMock)
|
||||
|
||||
reports := make([]*AWSMeterReport, 1)
|
||||
reports[0] = &AWSMeterReport{
|
||||
Dimension: model.AWS_METERING_DIMENSION_USAGE_HRS,
|
||||
Value: 2,
|
||||
Timestamp: startTime,
|
||||
}
|
||||
|
||||
// Define a mock struct to be used in your unit tests of myFunc.
|
||||
svc := &AWSMeterService{
|
||||
AwsDryRun: false,
|
||||
AwsProductCode: "12345",
|
||||
AwsMeteringSvc: &mockMarketplaceMeteringClient{},
|
||||
}
|
||||
|
||||
config := &model.Config{}
|
||||
config.SetDefaults()
|
||||
|
||||
awsmeter := &AwsMeter{
|
||||
store: &storeMock,
|
||||
service: svc,
|
||||
config: config,
|
||||
}
|
||||
|
||||
t.Run("Error in DB query", func(t *testing.T) {
|
||||
resultReports := awsmeter.GetUserCategoryUsage(dimensions, startTime, endTime)
|
||||
require.NotNil(t, resultReports)
|
||||
assert.Equal(t, 0, len(resultReports))
|
||||
err := awsmeter.ReportUserCategoryUsage(resultReports)
|
||||
require.Nil(t, err)
|
||||
})
|
||||
}
|
||||
Ссылка в новой задаче
Block a user