update test code

db/workerinfo_test.go

@@ -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()