Commit 53c21124 authored by luxq's avatar luxq

update test code

parent fcc1a079
......@@ -12,6 +12,7 @@ import (
"math/rand"
"mogo/types"
"strconv"
"sync"
"testing"
"time"
)
......@@ -189,6 +190,11 @@ func generateAHardware() *types.HardwareInfo {
generateAIdleGpu(0),
generateAIdleGpu(1),
generateAUsageGpu(2),
generateAUsageGpu(3),
generateAUsageGpu(4),
generateAUsageGpu(5),
generateAUsageGpu(6),
generateAUsageGpu(7),
},
RAM: generateARam(),
DISK: generateADisk(),
......@@ -205,6 +211,15 @@ func generateAInstallModel() *types.InstalledModel {
}
}
func generateAInstallModelWithId(mid string) *types.InstalledModel {
return &types.InstalledModel{
ModelID: mid,
DiskSize: 101,
InstalledTime: time.Now().Unix(),
LastRunTime: time.Now().Unix(),
}
}
func generateARunningModel() *types.RunningModel {
return &types.RunningModel{
ModelID: getRandId(maxModelId),
......@@ -217,20 +232,65 @@ func generateARunningModel() *types.RunningModel {
}
}
func generateARunningModelWithId(mid string) *types.RunningModel {
return &types.RunningModel{
ModelID: mid,
GpuSeq: rand.Intn(3),
GpuRAM: generateAGpuRam(),
StartedTime: time.Now().Unix(),
LastWorkTime: time.Now().Unix(),
TotalRunCount: rand.Intn(100),
WaitTime: rand.Intn(100),
}
}
func generateAModel() *types.ModelInfo {
m := &types.ModelInfo{
InstalledModels: make([]*types.InstalledModel, 0, 1000),
RunningModels: make([]*types.RunningModel, 0, 1000),
}
for i := 0; i < 500; i++ {
m.InstalledModels = append(m.InstalledModels, generateAInstallModel())
for i := 0; i < 100; i++ {
m.InstalledModels = append(m.InstalledModels, generateAInstallModelWithId(strconv.Itoa(i+1)))
if len(m.RunningModels) < 500 {
m.RunningModels = append(m.RunningModels, generateARunningModel())
m.RunningModels = append(m.RunningModels, generateARunningModelWithId(strconv.Itoa(i+1)))
}
}
return m
}
func TestDbWorker_InsertWorker(t *testing.T) {
client, err := ConnectMongoDB("mongodb://localhost:27017")
if err != nil {
t.Fatal(err)
}
db := NewDBWorker(client, database, collection)
defer db.client.Disconnect(context.Background())
total := 10000
count := 0
mux := sync.Mutex{}
wg := sync.WaitGroup{}
for i := 0; i < 5; i++ {
go func() {
defer wg.Done()
for {
worker := generateAWroker()
if _, err := db.InsertWorker(context.Background(), worker); err != nil {
t.Fatalf("insert worker failed with err:%s", err)
}
mux.Lock()
count++
mux.Unlock()
if count >= total {
return
}
}
}()
}
wg.Wait()
}
func BenchmarkGenerateWorker(b *testing.B) {
for i := 0; i < b.N; i++ {
generateAWroker()
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment