storj/satellite/attribution/db_test.go

554 lines
17 KiB
Go
Raw Normal View History

// Copyright (C) 2019 Storj Labs, Inc.
// See LICENSE for copying information.
package attribution_test
import (
"bytes"
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"storj.io/common/pb"
"storj.io/common/testcontext"
"storj.io/common/testrand"
"storj.io/common/uuid"
"storj.io/storj/satellite"
"storj.io/storj/satellite/accounting"
"storj.io/storj/satellite/attribution"
"storj.io/storj/satellite/orders"
"storj.io/storj/satellite/satellitedb/satellitedbtest"
)
const (
remoteSize int64 = 5368709120
inlineSize int64 = 1024
egressSize int64 = 2048
)
type AttributionTestData struct {
name string
userAgent []byte
projectID uuid.UUID
bucketName []byte
bucketID []byte
start time.Time
end time.Time
dataInterval time.Time
bwStart time.Time
hours int64
hoursOfData int
padding int
remoteSize int64
inlineSize int64
egressSize int64
dataCounter int
expectedByteHours float64
expectedSegmentHours float64
expectedObjectHours float64
expectedHours float64
expectedEgress int64
}
func (testData *AttributionTestData) init() {
testData.bucketID = []byte(testData.projectID.String() + "/" + string(testData.bucketName))
testData.hours = int64(testData.end.Sub(testData.start).Hours())
testData.hoursOfData = int(testData.hours) + (testData.padding * 2)
testData.dataInterval = time.Date(testData.start.Year(), testData.start.Month(), testData.start.Day(), -testData.padding, -1, 0, 0, testData.start.Location())
testData.bwStart = time.Date(testData.start.Year(), testData.start.Month(), testData.start.Day(), -testData.padding, 0, 0, 0, testData.start.Location())
}
func TestDB(t *testing.T) {
satellitedbtest.Run(t, func(ctx *testcontext.Context, t *testing.T, db satellite.DB) {
attributionDB := db.Attribution()
project1, project2 := testrand.UUID(), testrand.UUID()
agent1, agent2 := []byte("agent1"), []byte("agent2")
infos := []*attribution.Info{
{project1, []byte("alpha"), agent1, time.Time{}},
{project1, []byte("beta"), agent2, time.Time{}},
{project2, []byte("alpha"), agent2, time.Time{}},
{project2, []byte("beta"), agent1, time.Time{}},
}
for _, info := range infos {
got, err := attributionDB.Insert(ctx, info)
require.NoError(t, err)
got.CreatedAt = time.Time{}
assert.Equal(t, info, got)
}
for _, info := range infos {
got, err := attributionDB.Get(ctx, info.ProjectID, info.BucketName)
require.NoError(t, err)
assert.Equal(t, info.UserAgent, got.UserAgent)
}
})
}
func TestQueryAttribution(t *testing.T) {
satellitedbtest.Run(t, func(ctx *testcontext.Context, t *testing.T, db satellite.DB) {
now := time.Now()
projectID := testrand.UUID()
userAgent := []byte("agent1")
alphaBucket := []byte("alpha")
betaBucket := []byte("beta")
testData := []AttributionTestData{
{
name: "new partnerID, userAgent, projectID, alpha",
userAgent: []byte("agent2"),
projectID: projectID,
bucketName: alphaBucket,
remoteSize: remoteSize,
inlineSize: inlineSize,
egressSize: egressSize,
start: time.Date(now.Year(), now.Month(), now.Day(), 0, 0, 0, 0, now.Location()),
end: time.Date(now.Year(), now.Month(), now.Day()+1, 0, 0, 0, 0, now.Location()),
padding: 2,
},
{
name: "partnerID, userAgent, new projectID, alpha",
userAgent: userAgent,
projectID: testrand.UUID(),
bucketName: alphaBucket,
remoteSize: remoteSize / 2,
inlineSize: inlineSize / 2,
egressSize: egressSize / 2,
start: time.Date(now.Year(), now.Month(), now.Day(), 0, 0, 0, 0, now.Location()),
end: time.Date(now.Year(), now.Month(), now.Day()+1, 0, 0, 0, 0, now.Location()),
padding: 2,
},
{
name: "new partnerID, userAgent, projectID, beta",
userAgent: []byte("agent3"),
projectID: projectID,
bucketName: betaBucket,
remoteSize: remoteSize / 3,
inlineSize: inlineSize / 3,
egressSize: egressSize / 3,
start: time.Date(now.Year(), now.Month(), now.Day(), 0, 0, 0, 0, now.Location()),
end: time.Date(now.Year(), now.Month(), now.Day()+1, 0, 0, 0, 0, now.Location()),
padding: 2,
},
{
name: "partnerID, userAgent new projectID, beta",
userAgent: userAgent,
projectID: testrand.UUID(),
bucketName: betaBucket,
remoteSize: remoteSize / 4,
inlineSize: inlineSize / 4,
egressSize: egressSize / 4,
start: time.Date(now.Year(), now.Month(), now.Day(), 0, 0, 0, 0, now.Location()),
end: time.Date(now.Year(), now.Month(), now.Day()+1, 0, 0, 0, 0, now.Location()),
padding: 2,
}}
for _, td := range testData {
td := td
td.init()
info := attribution.Info{td.projectID, td.bucketName, td.userAgent, time.Time{}}
_, err := db.Attribution().Insert(ctx, &info)
require.NoError(t, err)
for i := 0; i < td.hoursOfData; i++ {
createData(ctx, t, db, &td)
}
verifyData(ctx, t, db.Attribution(), &td)
}
})
}
func TestQueryAllAttribution(t *testing.T) {
satellitedbtest.Run(t, func(ctx *testcontext.Context, t *testing.T, db satellite.DB) {
now := time.Now()
projectID := testrand.UUID()
userAgent := []byte("agent1")
alphaBucket := []byte("alpha")
betaBucket := []byte("beta")
testData := []AttributionTestData{
{
name: "new partnerID, userAgent, projectID, alpha",
userAgent: []byte("agent2"),
projectID: projectID,
bucketName: alphaBucket,
remoteSize: remoteSize,
inlineSize: inlineSize,
egressSize: egressSize,
start: time.Date(now.Year(), now.Month(), now.Day(), 0, 0, 0, 0, now.Location()),
end: time.Date(now.Year(), now.Month(), now.Day()+1, 0, 0, 0, 0, now.Location()),
padding: 2,
},
{
name: "partnerID, userAgent, new projectID, alpha",
userAgent: userAgent,
projectID: testrand.UUID(),
bucketName: alphaBucket,
remoteSize: remoteSize / 2,
inlineSize: inlineSize / 2,
egressSize: egressSize / 2,
start: time.Date(now.Year(), now.Month(), now.Day(), 0, 0, 0, 0, now.Location()),
end: time.Date(now.Year(), now.Month(), now.Day()+1, 0, 0, 0, 0, now.Location()),
padding: 2,
},
{
name: "new partnerID, userAgent, projectID, beta",
userAgent: []byte("agent3"),
projectID: projectID,
bucketName: betaBucket,
remoteSize: remoteSize / 3,
inlineSize: inlineSize / 3,
egressSize: egressSize / 3,
start: time.Date(now.Year(), now.Month(), now.Day(), 0, 0, 0, 0, now.Location()),
end: time.Date(now.Year(), now.Month(), now.Day()+1, 0, 0, 0, 0, now.Location()),
padding: 2,
},
{
name: "partnerID, userAgent new projectID, beta",
userAgent: userAgent,
projectID: testrand.UUID(),
bucketName: betaBucket,
remoteSize: remoteSize / 4,
inlineSize: inlineSize / 4,
egressSize: egressSize / 4,
start: time.Date(now.Year(), now.Month(), now.Day(), 0, 0, 0, 0, now.Location()),
end: time.Date(now.Year(), now.Month(), now.Day()+1, 0, 0, 0, 0, now.Location()),
padding: 2,
}}
for i, td := range testData {
td := td
td.init()
info := attribution.Info{td.projectID, td.bucketName, td.userAgent, time.Time{}}
_, err := db.Attribution().Insert(ctx, &info)
require.NoError(t, err)
for i := 0; i < td.hoursOfData; i++ {
if i < td.hoursOfData/2 {
createTallyData(ctx, t, db.ProjectAccounting(), &td)
} else {
createBWData(ctx, t, db.Orders(), &td)
}
}
testData[i] = td
}
verifyAllData(ctx, t, db.Attribution(), testData, false)
})
}
func TestQueryAllAttributionNoStorage(t *testing.T) {
satellitedbtest.Run(t, func(ctx *testcontext.Context, t *testing.T, db satellite.DB) {
now := time.Now()
projectID := testrand.UUID()
userAgent := []byte("agent1")
alphaBucket := []byte("alpha")
betaBucket := []byte("beta")
testData := []AttributionTestData{
{
name: "new partnerID, userAgent, projectID, alpha",
userAgent: []byte("agent2"),
projectID: projectID,
bucketName: alphaBucket,
remoteSize: 0,
inlineSize: 0,
egressSize: egressSize,
start: time.Date(now.Year(), now.Month(), now.Day(), 0, 0, 0, 0, now.Location()),
end: time.Date(now.Year(), now.Month(), now.Day()+1, 0, 0, 0, 0, now.Location()),
padding: 2,
},
{
name: "partnerID, userAgent, new projectID, alpha",
userAgent: userAgent,
projectID: testrand.UUID(),
bucketName: alphaBucket,
remoteSize: 0,
inlineSize: 0,
egressSize: egressSize / 2,
start: time.Date(now.Year(), now.Month(), now.Day(), 0, 0, 0, 0, now.Location()),
end: time.Date(now.Year(), now.Month(), now.Day()+1, 0, 0, 0, 0, now.Location()),
padding: 2,
},
{
name: "new partnerID, userAgent, projectID, beta",
userAgent: []byte("agent3"),
projectID: projectID,
bucketName: betaBucket,
remoteSize: 0,
inlineSize: 0,
egressSize: egressSize / 3,
start: time.Date(now.Year(), now.Month(), now.Day(), 0, 0, 0, 0, now.Location()),
end: time.Date(now.Year(), now.Month(), now.Day()+1, 0, 0, 0, 0, now.Location()),
padding: 2,
},
{
name: "partnerID, userAgent new projectID, beta",
userAgent: userAgent,
projectID: testrand.UUID(),
bucketName: betaBucket,
remoteSize: 0,
inlineSize: 0,
egressSize: egressSize / 4,
start: time.Date(now.Year(), now.Month(), now.Day(), 0, 0, 0, 0, now.Location()),
end: time.Date(now.Year(), now.Month(), now.Day()+1, 0, 0, 0, 0, now.Location()),
padding: 2,
}}
for i, td := range testData {
td := td
td.init()
info := attribution.Info{td.projectID, td.bucketName, td.userAgent, time.Time{}}
_, err := db.Attribution().Insert(ctx, &info)
require.NoError(t, err)
for i := 0; i < td.hoursOfData; i++ {
createBWData(ctx, t, db.Orders(), &td)
}
testData[i] = td
}
verifyAllData(ctx, t, db.Attribution(), testData, false)
})
}
func TestQueryAllAttributionNoBW(t *testing.T) {
satellitedbtest.Run(t, func(ctx *testcontext.Context, t *testing.T, db satellite.DB) {
now := time.Now()
projectID := testrand.UUID()
userAgent := []byte("agent1")
alphaBucket := []byte("alpha")
betaBucket := []byte("beta")
testData := []AttributionTestData{
{
name: "new partnerID, userAgent, projectID, alpha",
userAgent: []byte("agent2"),
projectID: projectID,
bucketName: alphaBucket,
remoteSize: remoteSize,
inlineSize: inlineSize,
egressSize: 0,
start: time.Date(now.Year(), now.Month(), now.Day(), 0, 0, 0, 0, now.Location()),
end: time.Date(now.Year(), now.Month(), now.Day()+1, 0, 0, 0, 0, now.Location()),
padding: 2,
},
{
name: "partnerID, userAgent, new projectID, alpha",
userAgent: userAgent,
projectID: testrand.UUID(),
bucketName: alphaBucket,
remoteSize: remoteSize / 2,
inlineSize: inlineSize / 2,
egressSize: 0,
start: time.Date(now.Year(), now.Month(), now.Day(), 0, 0, 0, 0, now.Location()),
end: time.Date(now.Year(), now.Month(), now.Day()+1, 0, 0, 0, 0, now.Location()),
padding: 2,
},
{
name: "new partnerID, userAgent, projectID, beta",
userAgent: []byte("agent3"),
projectID: projectID,
bucketName: betaBucket,
remoteSize: remoteSize / 3,
inlineSize: inlineSize / 3,
egressSize: 0,
start: time.Date(now.Year(), now.Month(), now.Day(), 0, 0, 0, 0, now.Location()),
end: time.Date(now.Year(), now.Month(), now.Day()+1, 0, 0, 0, 0, now.Location()),
padding: 2,
},
{
name: "partnerID, userAgent new projectID, beta",
userAgent: userAgent,
projectID: testrand.UUID(),
bucketName: betaBucket,
remoteSize: remoteSize / 4,
inlineSize: inlineSize / 4,
egressSize: 0,
start: time.Date(now.Year(), now.Month(), now.Day(), 0, 0, 0, 0, now.Location()),
end: time.Date(now.Year(), now.Month(), now.Day()+1, 0, 0, 0, 0, now.Location()),
padding: 2,
}}
for i, td := range testData {
td := td
td.init()
info := attribution.Info{td.projectID, td.bucketName, td.userAgent, time.Time{}}
_, err := db.Attribution().Insert(ctx, &info)
require.NoError(t, err)
for i := 0; i < td.hoursOfData; i++ {
createTallyData(ctx, t, db.ProjectAccounting(), &td)
}
testData[i] = td
}
verifyAllData(ctx, t, db.Attribution(), testData, true)
})
}
func verifyData(ctx *testcontext.Context, t *testing.T, attributionDB attribution.DB, testData *AttributionTestData) {
results, err := attributionDB.QueryAttribution(ctx, testData.userAgent, testData.start, testData.end)
require.NoError(t, err)
require.NotEqual(t, 0, len(results), "Results must not be empty.")
count := 0
for _, r := range results {
projectID, err := uuid.FromBytes(r.ProjectID)
require.NoError(t, err)
// The query returns results by partnerID, so we need to filter out by projectID
if projectID != testData.projectID {
continue
}
count++
assert.Equal(t, testData.userAgent, r.UserAgent, testData.name)
assert.Equal(t, testData.projectID[:], r.ProjectID, testData.name)
assert.Equal(t, testData.bucketName, r.BucketName, testData.name)
assert.Equal(t, testData.expectedByteHours, r.ByteHours, testData.name)
assert.Equal(t, testData.expectedEgress, r.EgressData, testData.name)
}
require.NotEqual(t, 0, count, "Results were returned, but did not match all of the the projectIDs.")
}
func verifyAllData(ctx *testcontext.Context, t *testing.T, attributionDB attribution.DB, testData []AttributionTestData, storageInEveryHour bool) {
results, err := attributionDB.QueryAllAttribution(ctx, testData[0].start, testData[0].end)
require.NoError(t, err)
require.NotEqual(t, 0, len(results), "Results must not be empty.")
count := 0
for _, tt := range testData {
for _, r := range results {
projectID, err := uuid.FromBytes(r.ProjectID)
require.NoError(t, err)
// The query returns results by partnerID, so we need to filter out by projectID
if projectID != tt.projectID || !bytes.Equal(r.BucketName, tt.bucketName) {
continue
}
count++
assert.Equal(t, tt.userAgent, r.UserAgent, tt.name)
assert.Equal(t, tt.projectID[:], r.ProjectID, tt.name)
assert.Equal(t, tt.bucketName, r.BucketName, tt.name)
assert.Equal(t, tt.expectedByteHours, r.ByteHours, tt.name)
assert.Equal(t, tt.expectedEgress, r.EgressData, tt.name)
if storageInEveryHour {
assert.Equal(t, tt.expectedHours, float64(tt.hours))
} else {
assert.Less(t, tt.expectedHours, float64(tt.hours))
}
}
}
require.Equal(t, len(testData), len(results))
require.NotEqual(t, 0, count, "Results were returned, but did not match all of the the projectIDs.")
}
func createData(ctx *testcontext.Context, t *testing.T, db satellite.DB, testData *AttributionTestData) {
createBWData(ctx, t, db.Orders(), testData)
createTallyData(ctx, t, db.ProjectAccounting(), testData)
}
func createBWData(ctx *testcontext.Context, t *testing.T, orderDB orders.DB, testData *AttributionTestData) {
err := orderDB.UpdateBucketBandwidthSettle(ctx, testData.projectID, testData.bucketName, pb.PieceAction_GET, testData.egressSize, 0, testData.bwStart)
require.NoError(t, err)
// Only GET should be counted. So this should not effect results
err = orderDB.UpdateBucketBandwidthSettle(ctx, testData.projectID, testData.bucketName, pb.PieceAction_GET_AUDIT, testData.egressSize, 0, testData.bwStart)
require.NoError(t, err)
if (testData.bwStart.After(testData.start) || testData.bwStart.Equal(testData.start)) && testData.bwStart.Before(testData.end) {
testData.expectedEgress += testData.egressSize
}
testData.bwStart = testData.bwStart.Add(time.Hour)
}
func createTallyData(ctx *testcontext.Context, t *testing.T, projectAccoutingDB accounting.ProjectAccounting, testData *AttributionTestData) {
testData.dataCounter++
testData.dataInterval = testData.dataInterval.Add(time.Minute * 30)
_, err := insertTallyData(ctx, projectAccoutingDB, testData.projectID, testData.bucketName, testData.remoteSize*int64(testData.dataCounter), testData.inlineSize*int64(testData.dataCounter), testData.dataInterval)
require.NoError(t, err)
// I think inserting into tallies twice is testing that only the latest row in the hour is used in the VA calculation.
testData.dataCounter++
testData.dataInterval = testData.dataInterval.Add(time.Minute * 30)
tally, err := insertTallyData(ctx, projectAccoutingDB, testData.projectID, testData.bucketName, testData.remoteSize*int64(testData.dataCounter), testData.inlineSize*int64(testData.dataCounter), testData.dataInterval)
require.NoError(t, err)
if (testData.dataInterval.After(testData.start) || testData.dataInterval.Equal(testData.start)) && testData.dataInterval.Before(testData.end) {
testData.expectedByteHours += float64(tally.TotalBytes)
testData.expectedHours++
testData.expectedSegmentHours += float64(tally.TotalSegmentCount)
testData.expectedObjectHours += float64(tally.ObjectCount)
}
}
func insertTallyData(ctx *testcontext.Context, projectAccoutingDB accounting.ProjectAccounting, projectID uuid.UUID, bucket []byte, remote int64, inline int64, dataIntervalStart time.Time) (_ accounting.BucketStorageTally, err error) {
tally := accounting.BucketStorageTally{
BucketName: string(bucket),
ProjectID: projectID,
IntervalStart: dataIntervalStart,
ObjectCount: 1,
TotalSegmentCount: 1,
TotalBytes: inline + remote,
MetadataSize: 0,
}
err = projectAccoutingDB.CreateStorageTally(ctx, tally)
if err != nil {
return accounting.BucketStorageTally{}, err
}
return tally, nil
}