update test

operator/workerInstalledInfo_test.go

@@ -10,7 +10,7 @@ import (
 )
 
 func TestWorkerInstalledOperator_Insert(t *testing.T) {
-	client, err := ConnectMongoDB("mongodb://localhost:27017", "admin", "admin")
+	client, err := ConnectMongoDB()
 	if err != nil {
 		log.Fatal(err)
 	}
@@ -53,7 +53,7 @@ func TestWorkerInstalledOperator_Insert(t *testing.T) {
 
 func BenchmarkWorkerInstalledOperator_UpdateGpuFree(b *testing.B) {
 
-	client, err := ConnectMongoDB("mongodb://localhost:27017", "admin", "admin")
+	client, err := ConnectMongoDB()
 	if err != nil {
 		log.Fatal(err)
 	}
@@ -71,14 +71,15 @@ func BenchmarkWorkerInstalledOperator_UpdateGpuFree(b *testing.B) {
 		idx := rand.Intn(len(allWorker))
 		workerid := allWorker[idx]
 		gpufree := rand.Intn(100) + 12000
-		if err := db.UpdateGpuFree(context.Background(), workerid, int64(gpufree)); err != nil {
+		gpuseq := rand.Intn(3)
+		if err := db.UpdateGpuFree(context.Background(), workerid, int64(gpufree), gpuseq); err != nil {
 			panic(fmt.Sprintf("update worker failed with err:%s", err))
 		}
 	}
 }
 
 func BenchmarkWorkerInstalledOperator_UpdateGpuFree_Parallel(b *testing.B) {
-	client, err := ConnectMongoDB("mongodb://localhost:27017", "admin", "admin")
+	client, err := ConnectMongoDB()
 	if err != nil {
 		log.Fatal(err)
 	}
@@ -98,53 +99,16 @@ func BenchmarkWorkerInstalledOperator_UpdateGpuFree_Parallel(b *testing.B) {
 			idx := rand.Intn(len(allWorker))
 			workerid := allWorker[idx]
 			gpufree := rand.Intn(100) + 12000
-			if err := db.UpdateGpuFree(context.Background(), workerid, int64(gpufree)); err != nil {
+			gpuseq := rand.Intn(3)
+			if err := db.UpdateGpuFree(context.Background(), workerid, int64(gpufree), gpuseq); err != nil {
 				panic(fmt.Sprintf("update worker failed with err:%s", err))
 			}
 		}
 	})
 }
 
-func BenchmarkWorkerInstalledOperator_FindWorkerByModelId(b *testing.B) {
-	client, err := ConnectMongoDB("mongodb://localhost:27017", "admin", "admin")
-	if err != nil {
-		log.Fatal(err)
-	}
-	db := NewDBWorkerInstalled(client, database)
-	defer db.client.Disconnect(context.Background())
-	b.ResetTimer()
-	for i := 0; i < b.N; i++ {
-		id := rand.Intn(maxModelId)
-		_, err := db.FindWorkerByModelId(context.Background(), id, 10)
-		if err != nil {
-			panic(fmt.Sprintf("find worker failed with err:%s", err))
-		}
-	}
-
-}
-
-func BenchmarkWorkerInstalledOperator_FindWorkerByModelId_Parallel(b *testing.B) {
-	client, err := ConnectMongoDB("mongodb://localhost:27017", "admin", "admin")
-	if err != nil {
-		log.Fatal(err)
-	}
-	db := NewDBWorkerInstalled(client, database)
-	defer db.client.Disconnect(context.Background())
-	b.ResetTimer()
-	b.RunParallel(func(pb *testing.PB) {
-		for pb.Next() {
-			id := rand.Intn(maxModelId)
-			_, err := db.FindWorkerByModelId(context.Background(), id, 10)
-			if err != nil {
-				panic(fmt.Sprintf("find worker failed with err:%s", err))
-			}
-		}
-	})
-
-}
-
 func BenchmarkWorkerInstalledOperator_FindWorkerByModelIdAndGpuMem(b *testing.B) {
-	client, err := ConnectMongoDB("mongodb://localhost:27017", "admin", "admin")
+	client, err := ConnectMongoDB()
 	if err != nil {
 		log.Fatal(err)
 	}
@@ -162,7 +126,7 @@ func BenchmarkWorkerInstalledOperator_FindWorkerByModelIdAndGpuMem(b *testing.B)
 }
 
 func BenchmarkWorkerInstalledOperator_FindWorkerByModelIdAndGpuMem_Parallel(b *testing.B) {
-	client, err := ConnectMongoDB("mongodb://localhost:27017", "admin", "admin")
+	client, err := ConnectMongoDB()
 	if err != nil {
 		log.Fatal(err)
 	}
@@ -183,7 +147,7 @@ func BenchmarkWorkerInstalledOperator_FindWorkerByModelIdAndGpuMem_Parallel(b *t
 }
 
 func TestWorkerInstalledOperator_UpdateGpuFree(t *testing.T) {
-	client, err := ConnectMongoDB("mongodb://localhost:27017", "admin", "admin")
+	client, err := ConnectMongoDB()
 	if err != nil {
 		log.Fatal(err)
 	}
@@ -191,7 +155,8 @@ func TestWorkerInstalledOperator_UpdateGpuFree(t *testing.T) {
 	defer db.client.Disconnect(context.Background())
 	workerid := "ebb6a2df-9e91-4b3c-ace8-748aaeca718c"
 	gpufree := 102222
-	if err := db.UpdateGpuFree(context.Background(), workerid, int64(gpufree)); err != nil {
+	gpuseq := 3
+	if err := db.UpdateGpuFree(context.Background(), workerid, int64(gpufree), gpuseq); err != nil {
 		panic(fmt.Sprintf("update worker failed with err:%s", err))
 	}
 }

operator/workerRunningInfo.go

@@ -3,6 +3,7 @@ package operator
 import (
 	"context"
 	"go.mongodb.org/mongo-driver/bson"
+	"go.mongodb.org/mongo-driver/bson/primitive"
 	"go.mongodb.org/mongo-driver/mongo"
 	"go.mongodb.org/mongo-driver/mongo/options"
 )
@@ -50,6 +51,13 @@ func (d *WorkerRunningOperator) UpdateExecTime(ctx context.Context, workerid str
 	return err
 }
 
+func (d *WorkerRunningOperator) Get(ctx context.Context, id string) (*WorkerRunningInfo, error) {
+	var worker WorkerRunningInfo
+	oid, _ := primitive.ObjectIDFromHex(id)
+	err := d.col.FindOne(ctx, bson.M{"_id": oid}).Decode(&worker)
+	return &worker, err
+}
+
 func (d *WorkerRunningOperator) DeleteByWorkerId(ctx context.Context, workerid string) (int, error) {
 	res, err := d.col.DeleteMany(ctx, bson.M{"worker_id": workerid})
 	return int(res.DeletedCount), err

operator/workerRunningInfo_test.go

+package operator
+
+import (
+	"context"
+	"fmt"
+	"github.com/google/uuid"
+	"go.mongodb.org/mongo-driver/bson/primitive"
+	"log"
+	"testing"
+)
+
+func TestWorkerRunningOperator_Insert(t *testing.T) {
+	client, err := ConnectMongoDB()
+	if err != nil {
+		log.Fatal(err)
+	}
+	db := NewDBWorkerRunning(client, database)
+	defer db.client.Disconnect(context.Background())
+	w := &WorkerRunningInfo{
+		WorkerId: uuid.NewString(),
+		ModelId:  1,
+		ExecTime: 100,
+	}
+	is, err := db.Insert(context.Background(), w)
+	if err != nil {
+		t.Errorf("insert worker failed with err:%s", err)
+	}
+	id := is.InsertedID.(primitive.ObjectID)
+
+	worker, _ := db.Get(context.Background(), id.Hex())
+	if worker.WorkerId != w.WorkerId {
+		t.Errorf("insert worker failed with err:%s", err)
+	}
+}
+
+func TestWorkerRunningOperator_InsertMany(t *testing.T) {
+	client, err := ConnectMongoDB()
+	if err != nil {
+		log.Fatal(err)
+	}
+	db := NewDBWorkerRunning(client, database)
+	defer db.client.Disconnect(context.Background())
+	ws := make([]*WorkerRunningInfo, 0)
+	for i := 0; i < 10; i++ {
+		w := &WorkerRunningInfo{
+			WorkerId: uuid.NewString(),
+			ModelId:  i,
+			ExecTime: 100,
+		}
+		ws = append(ws, w)
+	}
+	if res, err := db.InsertMany(context.Background(), ws); err != nil {
+		t.Errorf("insert worker failed with err:%s", err)
+	} else {
+		if len(res.InsertedIDs) != len(ws) {
+			t.Errorf("insert worker failed with err:%s", err)
+		}
+	}
+}
+
+func TestWorkerRunningOperator_DeleteMany(t *testing.T) {
+	client, err := ConnectMongoDB()
+	if err != nil {
+		log.Fatal(err)
+	}
+	db := NewDBWorkerRunning(client, database)
+	defer db.client.Disconnect(context.Background())
+	ws := make([]*WorkerRunningInfo, 0)
+	wid := uuid.NewString()
+	for i := 0; i < 10; i++ {
+		w := &WorkerRunningInfo{
+			WorkerId: wid,
+			ModelId:  i,
+			ExecTime: 100,
+		}
+		ws = append(ws, w)
+	}
+	if _, err := db.InsertMany(context.Background(), ws); err != nil {
+		t.Errorf("insert worker failed with err:%s", err)
+	}
+	deleteId := []int{1, 2, 3}
+	if n, err := db.DeleteMany(context.Background(), wid, deleteId); err != nil {
+		t.Errorf("delete worker failed with err:%s", err)
+	} else {
+		if n != len(deleteId) {
+			t.Errorf("delete worker failed with err:%s", err)
+		}
+	}
+}
+
+func TestWorkerRunningOperator_DeleteByWorkerId(t *testing.T) {
+	client, err := ConnectMongoDB()
+	if err != nil {
+		log.Fatal(err)
+	}
+	db := NewDBWorkerRunning(client, database)
+	defer db.client.Disconnect(context.Background())
+	ws := make([]*WorkerRunningInfo, 0)
+	wid := uuid.NewString()
+	for i := 0; i < 10; i++ {
+		w := &WorkerRunningInfo{
+			WorkerId: wid,
+			ModelId:  i,
+			ExecTime: 100,
+		}
+		ws = append(ws, w)
+	}
+	if _, err := db.InsertMany(context.Background(), ws); err != nil {
+		t.Errorf("insert worker failed with err:%s", err)
+	}
+	if n, err := db.DeleteByWorkerId(context.Background(), wid); err != nil {
+		t.Errorf("delete worker failed with err:%s", err)
+	} else {
+		if n != len(ws) {
+			t.Errorf("delete worker failed with err:%s", err)
+		}
+	}
+}
+
+func TestWorkerRunningOperator_UpdateExecTime(t *testing.T) {
+	client, err := ConnectMongoDB()
+	if err != nil {
+		log.Fatal(err)
+	}
+	db := NewDBWorkerRunning(client, database)
+	defer db.client.Disconnect(context.Background())
+	w := &WorkerRunningInfo{
+		WorkerId: uuid.NewString(),
+		ModelId:  1,
+		ExecTime: 100,
+	}
+	is, err := db.Insert(context.Background(), w)
+	if err != nil {
+		t.Errorf("insert worker failed with err:%s", err)
+	}
+
+	if err := db.UpdateExecTime(context.Background(), w.WorkerId, w.ModelId, 200); err != nil {
+		t.Errorf("update worker failed with err:%s", err)
+	}
+	id := is.InsertedID.(primitive.ObjectID)
+	info, err := db.Get(context.Background(), id.Hex())
+	if err != nil {
+		t.Errorf("update worker failed with err:%s", err)
+	}
+	if info.ExecTime != 200 {
+		t.Errorf("update worker failed with err:%s", err)
+	}
+}
+
+func BenchmarkDbWorkerRunning_FindWorkerByModelId(b *testing.B) {
+	client, err := ConnectMongoDB()
+	if err != nil {
+		log.Fatal(err)
+	}
+	db := NewDBWorkerRunning(client, database)
+	defer db.client.Disconnect(context.Background())
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		modelId := getRandIdInt(maxModelId)
+		if w, err := db.FindWorkerByModelId(context.Background(), modelId, 10); err != nil {
+			panic(fmt.Sprintf("find worker failed with err:%s", err))
+		} else if len(w) == 0 {
+			b.Logf("FindWorkerByModelId find %d with id %d\n", len(w), modelId)
+		}
+	}
+}
+
+func BenchmarkDbWorkerRunning_FindWorkerByModelId_Parallel(b *testing.B) {
+	client, err := ConnectMongoDB()
+	if err != nil {
+		log.Fatal(err)
+	}
+	db := NewDBWorkerRunning(client, database)
+	defer db.client.Disconnect(context.Background())
+	b.RunParallel(func(pb *testing.PB) {
+		for pb.Next() {
+			modelId := getRandIdInt(maxModelId)
+			if w, err := db.FindWorkerByModelId(context.Background(), modelId, 10); err != nil {
+				panic(fmt.Sprintf("find worker failed with err:%s", err))
+			} else if len(w) == 0 {
+				b.Logf("FindWorkerByModelId find %d with id %d\n", len(w), modelId)
+			}
+		}
+	})
+}

operator/workerinfo_test.go

@@ -27,7 +27,10 @@ var (
 	workerInstalledCollection = "worker_installed"
 )
 
-func ConnectMongoDB(uri string, username, passwd string) (*mongo.Client, error) {
+func ConnectMongoDB() (*mongo.Client, error) {
+	uri := "mongodb://localhost:27017"
+	username := "admin"
+	passwd := "admin"
 	ctx, _ := context.WithTimeout(context.Background(), 10*time.Second)
 	client, err := mongo.Connect(ctx, options.Client().ApplyURI(uri).SetAuth(options.Credential{
 		Username: username,
@@ -41,7 +44,7 @@ func ConnectMongoDB(uri string, username, passwd string) (*mongo.Client, error)
 }
 
 //func init() {
-//	client, err := ConnectMongoDB("mongodb://localhost:27017", "admin", "admin")
+//	client, err := ConnectMongoDB()
 //	if err != nil {
 //		log.Fatal(err)
 //	}
@@ -257,7 +260,7 @@ func BenchmarkGenerateWorker(b *testing.B) {
 }
 
 func BenchmarkDbWorker_InsertWorker(b *testing.B) {
-	client, err := ConnectMongoDB("mongodb://localhost:27017", "admin", "admin")
+	client, err := ConnectMongoDB()
 	if err != nil {
 		log.Fatal(err)
 	}
@@ -275,7 +278,7 @@ func BenchmarkDbWorker_InsertWorker(b *testing.B) {
 }
 
 func BenchmarkDbWorker_InsertWorker_Parallel(b *testing.B) {
-	client, err := ConnectMongoDB("mongodb://localhost:27017", "admin", "admin")
+	client, err := ConnectMongoDB()
 	if err != nil {
 		log.Fatal(err)
 	}
@@ -292,7 +295,7 @@ func BenchmarkDbWorker_InsertWorker_Parallel(b *testing.B) {
 }
 
 func BenchmarkDbWorker_UpdateHardware(b *testing.B) {
-	client, err := ConnectMongoDB("mongodb://localhost:27017", "admin", "admin")
+	client, err := ConnectMongoDB()
 	if err != nil {
 		log.Fatal(err)
 	}
@@ -314,7 +317,7 @@ func BenchmarkDbWorker_UpdateHardware(b *testing.B) {
 }
 
 func BenchmarkDbWorker_UpdateHardware_Parallel(b *testing.B) {
-	client, err := ConnectMongoDB("mongodb://localhost:27017", "admin", "admin")
+	client, err := ConnectMongoDB()
 	if err != nil {
 		log.Fatal(err)
 	}
@@ -332,81 +335,8 @@ func BenchmarkDbWorker_UpdateHardware_Parallel(b *testing.B) {
 	})
 }
 
-func BenchmarkDbWorker_UpdateModel(b *testing.B) {
-	client, err := ConnectMongoDB("mongodb://localhost:27017", "admin", "admin")
-	if err != nil {
-		log.Fatal(err)
-	}
-	b.ResetTimer()
-
-	db := NewDBWorker(client, database)
-	defer db.client.Disconnect(context.Background())
-	for i := 0; i < b.N; i++ {
-		idx := rand.Intn(len(idlist))
-		nresource := generateAModel()
-		if err := db.UpdateModel(context.Background(), idlist[idx], nresource); err != nil {
-			panic(fmt.Sprintf("update worker failed with err:%s", err))
-		}
-	}
-}
-
-func BenchmarkDbWorker_UpdateModel_Parallel(b *testing.B) {
-	client, err := ConnectMongoDB("mongodb://localhost:27017", "admin", "admin")
-	if err != nil {
-		log.Fatal(err)
-	}
-
-	db := NewDBWorker(client, database)
-	defer db.client.Disconnect(context.Background())
-	b.RunParallel(func(pb *testing.PB) {
-		for pb.Next() {
-			idx := rand.Intn(len(idlist))
-			nresource := generateAModel()
-			if err := db.UpdateModel(context.Background(), idlist[idx], nresource); err != nil {
-				panic(fmt.Sprintf("update worker failed with err:%s", err))
-			}
-		}
-	})
-}
-
-func BenchmarkDbWorker_UpdateNodeInfo(b *testing.B) {
-	client, err := ConnectMongoDB("mongodb://localhost:27017", "admin", "admin")
-	if err != nil {
-		log.Fatal(err)
-	}
-	b.ResetTimer()
-	nnodeinfo := generateANodeInfo()
-
-	db := NewDBWorker(client, database)
-	defer db.client.Disconnect(context.Background())
-	for i := 0; i < b.N; i++ {
-		idx := rand.Intn(len(idlist))
-		if err := db.UpdateNodeInfo(context.Background(), idlist[idx], nnodeinfo); err != nil {
-			panic(fmt.Sprintf("update worker failed with err:%s", err))
-		}
-	}
-}
-
-func BenchmarkDbWorker_UpdateNodeInfo_Parallel(b *testing.B) {
-	client, err := ConnectMongoDB("mongodb://localhost:27017", "admin", "admin")
-	if err != nil {
-		log.Fatal(err)
-	}
-	nnodeinfo := generateANodeInfo()
-	db := NewDBWorker(client, database)
-	defer db.client.Disconnect(context.Background())
-	b.RunParallel(func(pb *testing.PB) {
-		for pb.Next() {
-			idx := rand.Intn(len(idlist))
-			if err := db.UpdateNodeInfo(context.Background(), idlist[idx], nnodeinfo); err != nil {
-				panic(fmt.Sprintf("update worker failed with err:%s", err))
-			}
-		}
-	})
-}
-
 func BenchmarkDbWorker_FindWorkerByInstallModelAndSortByGpuRam(b *testing.B) {
-	client, err := ConnectMongoDB("mongodb://localhost:27017", "admin", "admin")
+	client, err := ConnectMongoDB()
 	if err != nil {
 		log.Fatal(err)
 	}
@@ -426,7 +356,7 @@ func BenchmarkDbWorker_FindWorkerByInstallModelAndSortByGpuRam(b *testing.B) {
 }
 
 func BenchmarkDbWorker_FindWorkerByInstallModelAndSortByGpuRam_Parallel(b *testing.B) {
-	client, err := ConnectMongoDB("mongodb://localhost:27017", "admin", "admin")
+	client, err := ConnectMongoDB()
 	if err != nil {
 		log.Fatal(err)
 	}
@@ -447,7 +377,7 @@ func BenchmarkDbWorker_FindWorkerByInstallModelAndSortByGpuRam_Parallel(b *testi
 }
 
 func BenchmarkDbWorker_FindWorkerByRunningModelAndSortByWaitTime(b *testing.B) {
-	client, err := ConnectMongoDB("mongodb://localhost:27017", "admin", "admin")
+	client, err := ConnectMongoDB()
 	if err != nil {
 		log.Fatal(err)
 	}
@@ -465,7 +395,7 @@ func BenchmarkDbWorker_FindWorkerByRunningModelAndSortByWaitTime(b *testing.B) {
 }
 
 func BenchmarkDbWorker_FindWorkerByRunningModelAndSortByWaitTime_Parallel(b *testing.B) {
-	client, err := ConnectMongoDB("mongodb://localhost:27017", "admin", "admin")
+	client, err := ConnectMongoDB()
 	if err != nil {
 		log.Fatal(err)
 	}
@@ -482,40 +412,3 @@ func BenchmarkDbWorker_FindWorkerByRunningModelAndSortByWaitTime_Parallel(b *tes
 		}
 	})
 }
-
-func BenchmarkDbWorkerRunning_FindWorkerByModelId(b *testing.B) {
-	client, err := ConnectMongoDB("mongodb://localhost:27017", "admin", "admin")
-	if err != nil {
-		log.Fatal(err)
-	}
-	db := NewDBWorkerRunning(client, database)
-	defer db.client.Disconnect(context.Background())
-	b.ResetTimer()
-	for i := 0; i < b.N; i++ {
-		modelId := getRandIdInt(maxModelId)
-		if w, err := db.FindWorkerByModelId(context.Background(), modelId, 10); err != nil {
-			panic(fmt.Sprintf("find worker failed with err:%s", err))
-		} else if len(w) == 0 {
-			b.Logf("FindWorkerByModelId find %d with id %d\n", len(w), modelId)
-		}
-	}
-}
-
-func BenchmarkDbWorkerRunning_FindWorkerByModelId_Parallel(b *testing.B) {
-	client, err := ConnectMongoDB("mongodb://localhost:27017", "admin", "admin")
-	if err != nil {
-		log.Fatal(err)
-	}
-	db := NewDBWorkerRunning(client, database)
-	defer db.client.Disconnect(context.Background())
-	b.RunParallel(func(pb *testing.PB) {
-		for pb.Next() {
-			modelId := getRandIdInt(maxModelId)
-			if w, err := db.FindWorkerByModelId(context.Background(), modelId, 10); err != nil {
-				panic(fmt.Sprintf("find worker failed with err:%s", err))
-			} else if len(w) == 0 {
-				b.Logf("FindWorkerByModelId find %d with id %d\n", len(w), modelId)
-			}
-		}
-	})
-}