mirror of https://github.com/harness/drone.git
114 lines
2.8 KiB
Go
114 lines
2.8 KiB
Go
// Copyright 2023 Harness, Inc.
|
|
//
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
// you may not use this file except in compliance with the License.
|
|
// You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
// See the License for the specific language governing permissions and
|
|
// limitations under the License.
|
|
|
|
package usage
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"sync"
|
|
"sync/atomic"
|
|
"testing"
|
|
|
|
"github.com/harness/gitness/types"
|
|
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
func TestMediator_basic(t *testing.T) {
|
|
space := &types.Space{
|
|
ID: 1,
|
|
Identifier: "space",
|
|
}
|
|
spaceMock := &SpaceStoreMock{
|
|
FindByRefFn: func(context.Context, string) (*types.Space, error) {
|
|
return space, nil
|
|
},
|
|
FindByIDsFn: func(context.Context, ...int64) ([]*types.Space, error) {
|
|
return []*types.Space{space}, nil
|
|
},
|
|
}
|
|
initialBandwidth := int64(1024)
|
|
initialStorage := int64(1024)
|
|
bandwidth := atomic.Int64{}
|
|
storage := atomic.Int64{}
|
|
counter := atomic.Int64{}
|
|
usageMock := &MetricsMock{
|
|
UpsertOptimisticFn: func(_ context.Context, in *types.UsageMetric) error {
|
|
if in.RootSpaceID != space.ID {
|
|
return fmt.Errorf("expected root space id to be %d, got %d", space.ID, in.RootSpaceID)
|
|
}
|
|
bandwidth.Add(in.Bandwidth)
|
|
storage.Add(in.Storage)
|
|
counter.Add(1)
|
|
return nil
|
|
},
|
|
GetMetricsFn: func(
|
|
context.Context,
|
|
int64, // spaceID
|
|
int64, // startDate
|
|
int64, // endDate
|
|
) (*types.UsageMetric, error) {
|
|
return &types.UsageMetric{
|
|
Bandwidth: bandwidth.Load(),
|
|
Storage: storage.Load(),
|
|
}, nil
|
|
},
|
|
ListFn: func(context.Context, int64, int64) ([]types.UsageMetric, error) {
|
|
bandwidth.Add(initialBandwidth)
|
|
storage.Add(initialStorage)
|
|
return []types.UsageMetric{
|
|
{
|
|
RootSpaceID: space.ID,
|
|
Bandwidth: initialBandwidth,
|
|
Storage: initialStorage,
|
|
},
|
|
}, nil
|
|
},
|
|
}
|
|
|
|
numRoutines := 10
|
|
defaultSize := 512
|
|
mediator := NewMediator(
|
|
context.Background(),
|
|
spaceMock,
|
|
usageMock,
|
|
Config{
|
|
ChunkSize: 1024,
|
|
MaxWorkers: 10,
|
|
},
|
|
)
|
|
wg := sync.WaitGroup{}
|
|
for range numRoutines {
|
|
wg.Add(1)
|
|
go func() {
|
|
defer wg.Done()
|
|
_ = mediator.Send(context.Background(), Metric{
|
|
SpaceRef: space.Identifier,
|
|
Size: Size{
|
|
Bandwidth: int64(defaultSize),
|
|
Storage: int64(defaultSize),
|
|
},
|
|
})
|
|
}()
|
|
}
|
|
|
|
wg.Wait()
|
|
mediator.Wait()
|
|
|
|
require.Equal(t, int64(numRoutines*defaultSize/int(mediator.config.ChunkSize)), counter.Load())
|
|
require.Equal(t, initialBandwidth+int64(numRoutines*defaultSize), bandwidth.Load())
|
|
require.Equal(t, initialStorage+int64(numRoutines*defaultSize), storage.Load())
|
|
}
|