github.com/masterhung0112/hk_server/v5@v5.0.0-20220302090640-ec71aef15e1c/services/awsmeter/awsmeter_test.go (about)

     1  // Copyright (c) 2015-present Mattermost, Inc. All Rights Reserved.
     2  // See LICENSE.txt for license information.
     3  
     4  package awsmeter
     5  
     6  import (
     7  	"errors"
     8  	"testing"
     9  	"time"
    10  
    11  	"github.com/aws/aws-sdk-go/service/marketplacemetering"
    12  	"github.com/aws/aws-sdk-go/service/marketplacemetering/marketplacemeteringiface"
    13  	"github.com/stretchr/testify/assert"
    14  	"github.com/stretchr/testify/require"
    15  
    16  	"github.com/masterhung0112/hk_server/v5/model"
    17  	"github.com/masterhung0112/hk_server/v5/plugin/plugintest/mock"
    18  	"github.com/masterhung0112/hk_server/v5/store/storetest/mocks"
    19  )
    20  
    21  type mockMarketplaceMeteringClient struct {
    22  	marketplacemeteringiface.MarketplaceMeteringAPI
    23  }
    24  
    25  func (m *mockMarketplaceMeteringClient) MeterUsage(input *marketplacemetering.MeterUsageInput) (*marketplacemetering.MeterUsageOutput, error) {
    26  	return &marketplacemetering.MeterUsageOutput{
    27  		MeteringRecordId: String("1"),
    28  	}, nil
    29  }
    30  
    31  type mockMarketplaceMeteringClientWithError struct {
    32  	marketplacemeteringiface.MarketplaceMeteringAPI
    33  }
    34  
    35  func (m *mockMarketplaceMeteringClientWithError) MeterUsage(input *marketplacemetering.MeterUsageInput) (*marketplacemetering.MeterUsageOutput, error) {
    36  	return nil, errors.New("error")
    37  }
    38  
    39  func String(i string) *string {
    40  	return &i
    41  }
    42  func TestAwsMeterUsage(t *testing.T) {
    43  	startTime := time.Now()
    44  	endTime := time.Now()
    45  	dimensions := []string{model.AWS_METERING_DIMENSION_USAGE_HRS}
    46  
    47  	userStoreMock := mocks.UserStore{}
    48  	userStoreMock.On("AnalyticsActiveCountForPeriod", model.GetMillisForTime(startTime), model.GetMillisForTime(endTime), mock.AnythingOfType("model.UserCountOptions")).Return(int64(2), nil)
    49  
    50  	storeMock := mocks.Store{}
    51  	storeMock.On("User").Return(&userStoreMock)
    52  
    53  	reports := make([]*AWSMeterReport, 1)
    54  	reports[0] = &AWSMeterReport{
    55  		Dimension: model.AWS_METERING_DIMENSION_USAGE_HRS,
    56  		Value:     2,
    57  		Timestamp: startTime,
    58  	}
    59  
    60  	// Define a mock struct to be used in your unit tests of myFunc.
    61  	svc := &AWSMeterService{
    62  		AwsDryRun:      false,
    63  		AwsProductCode: "12345",
    64  		AwsMeteringSvc: &mockMarketplaceMeteringClient{},
    65  	}
    66  
    67  	config := &model.Config{}
    68  	config.SetDefaults()
    69  
    70  	awsmeter := &AwsMeter{
    71  		store:   &storeMock,
    72  		service: svc,
    73  		config:  config,
    74  	}
    75  
    76  	t.Run("Send report for one usage category", func(t *testing.T) {
    77  		resultReports := awsmeter.GetUserCategoryUsage(dimensions, startTime, endTime)
    78  		require.NotNil(t, resultReports)
    79  		assert.Equal(t, 1, len(resultReports))
    80  		assert.Equal(t, reports[0].Dimension, resultReports[0].Dimension)
    81  		assert.Equal(t, reports[0].Value, resultReports[0].Value)
    82  		assert.Equal(t, reports[0].Timestamp, resultReports[0].Timestamp)
    83  
    84  		err := awsmeter.ReportUserCategoryUsage(resultReports)
    85  		require.NoError(t, err)
    86  	})
    87  
    88  	t.Run("Error in AWS service call", func(t *testing.T) {
    89  		awsmeter.service.AwsMeteringSvc = &mockMarketplaceMeteringClientWithError{}
    90  		resultReports := awsmeter.GetUserCategoryUsage(dimensions, startTime, endTime)
    91  		require.NotNil(t, resultReports)
    92  		assert.Equal(t, 1, len(resultReports))
    93  		err := awsmeter.ReportUserCategoryUsage(resultReports)
    94  		require.Error(t, err)
    95  	})
    96  
    97  	t.Run("Invalid dimension", func(t *testing.T) {
    98  		awsmeter.service.AwsMeteringSvc = &mockMarketplaceMeteringClient{}
    99  		dimensions = []string{"invalid dimension"}
   100  		resultReports := awsmeter.GetUserCategoryUsage(dimensions, startTime, endTime)
   101  		require.NotNil(t, resultReports)
   102  		assert.Equal(t, 0, len(resultReports))
   103  		err := awsmeter.ReportUserCategoryUsage(resultReports)
   104  		require.NoError(t, err)
   105  	})
   106  }
   107  
   108  func TestAwsMeterUsageWithDBError(t *testing.T) {
   109  	startTime := time.Now()
   110  	endTime := time.Now()
   111  	dimensions := []string{model.AWS_METERING_DIMENSION_USAGE_HRS}
   112  
   113  	userStoreMock := mocks.UserStore{}
   114  	userStoreMock.On("AnalyticsActiveCountForPeriod", model.GetMillisForTime(startTime), model.GetMillisForTime(endTime), mock.AnythingOfType("model.UserCountOptions")).Return(int64(0), errors.New("error"))
   115  
   116  	storeMock := mocks.Store{}
   117  	storeMock.On("User").Return(&userStoreMock)
   118  
   119  	reports := make([]*AWSMeterReport, 1)
   120  	reports[0] = &AWSMeterReport{
   121  		Dimension: model.AWS_METERING_DIMENSION_USAGE_HRS,
   122  		Value:     2,
   123  		Timestamp: startTime,
   124  	}
   125  
   126  	// Define a mock struct to be used in your unit tests of myFunc.
   127  	svc := &AWSMeterService{
   128  		AwsDryRun:      false,
   129  		AwsProductCode: "12345",
   130  		AwsMeteringSvc: &mockMarketplaceMeteringClient{},
   131  	}
   132  
   133  	config := &model.Config{}
   134  	config.SetDefaults()
   135  
   136  	awsmeter := &AwsMeter{
   137  		store:   &storeMock,
   138  		service: svc,
   139  		config:  config,
   140  	}
   141  
   142  	t.Run("Error in DB query", func(t *testing.T) {
   143  		resultReports := awsmeter.GetUserCategoryUsage(dimensions, startTime, endTime)
   144  		require.NotNil(t, resultReports)
   145  		assert.Equal(t, 0, len(resultReports))
   146  		err := awsmeter.ReportUserCategoryUsage(resultReports)
   147  		require.NoError(t, err)
   148  	})
   149  }