update test code

operator/workerInstalledInfo.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"
 )
@@ -55,6 +56,18 @@ func (d *WorkerInstalledOperator) DeleteByWorkerId(ctx context.Context, workerid
 	return int(res.DeletedCount), err
 }
 
+func (d *WorkerInstalledOperator) Get(ctx context.Context, id string) (*WorkerInstalledInfo, error) {
+	var worker WorkerInstalledInfo
+	oid, _ := primitive.ObjectIDFromHex(id)
+	err := d.col.FindOne(ctx, bson.M{"_id": oid}).Decode(&worker)
+	return &worker, err
+}
+
+func (d *WorkerInstalledOperator) Clear() int64 {
+	res, _ := d.col.DeleteMany(context.Background(), bson.M{})
+	return res.DeletedCount
+}
+
 //
 //func (d *WorkerInstalledOperator) FindWorkerByModelId(ctx context.Context, modelId int, limit int) ([]*WorkerInstalledInfo, error) {
 //	// find all worker that at least one installed model's mode_id is equal modelId

operator/workerInstalledInfo_test.go

@@ -3,54 +3,273 @@ package operator
 import (
 	"context"
 	"fmt"
+	"github.com/google/uuid"
+	"go.mongodb.org/mongo-driver/bson/primitive"
 	"log"
 	"math/rand"
-	"strconv"
 	"testing"
 )
 
-func TestWorkerInstalledOperator_Insert(t *testing.T) {
+func TestWorkerInstalledOperator_InsertMany(t *testing.T) {
 	client, err := ConnectMongoDB()
 	if err != nil {
 		log.Fatal(err)
 	}
-	count := 1
-	infodb := NewDBWorker(client, database)
-	defer infodb.client.Disconnect(context.Background())
 	db := NewDBWorkerInstalled(client, database)
 	defer db.client.Disconnect(context.Background())
-	rundb := NewDBWorkerRunning(client, database)
-	defer rundb.client.Disconnect(context.Background())
-
-	for i := 0; i < count; i++ {
-		worker := generateAWroker()
-		infodb.InsertWorker(context.Background(), worker)
-		for _, installed := range worker.Models.InstalledModels {
-			id, _ := strconv.Atoi(installed.ModelID)
-			_, err := db.Insert(context.Background(), &WorkerInstalledInfo{
-				WorkerId: worker.WorkerId,
-				ModelId:  id,
+
+	ws := make([]*WorkerInstalledInfo, 0)
+	for i := 0; i < 10; i++ {
+		w := &WorkerInstalledInfo{
+			WorkerId: uuid.NewString(),
+			ModelId:  i,
 			GpuFree:  10011,
-			})
+		}
+		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 TestWorkerInstalledOperator_DeleteMany(t *testing.T) {
+	client, err := ConnectMongoDB()
+	if err != nil {
+		log.Fatal(err)
+	}
+	db := NewDBWorkerInstalled(client, database)
+	defer db.client.Disconnect(context.Background())
+
+	workerid := uuid.NewString()
+	// insert first
+	ws := make([]*WorkerInstalledInfo, 0)
+	for i := 0; i < 10; i++ {
+		w := &WorkerInstalledInfo{
+			WorkerId: workerid,
+			ModelId:  i,
+			GpuFree:  10011,
+		}
+		ws = append(ws, w)
+	}
+	if _, err := db.InsertMany(context.Background(), ws); err != nil {
+		t.Errorf("insert worker failed with err:%s", err)
+	}
+	models := []int{1, 2, 3, 4, 5}
+	n, err := db.DeleteMany(context.Background(), workerid, models)
 	if err != nil {
-				panic(fmt.Sprintf("insert worker failed with err:%s", err))
+		t.Errorf("delete worker failed with err:%s", err)
 	}
+	if n != len(models) {
+		t.Errorf("delete worker failed with err:%s", err)
 	}
+}
 
-		for _, running := range worker.Models.RunningModels {
-			id, _ := strconv.Atoi(running.ModelID)
-			_, err := rundb.Insert(context.Background(), &WorkerRunningInfo{
-				WorkerId: worker.WorkerId,
-				ModelId:  id,
-				ExecTime: 100,
-			})
+func TestWorkerInstalledOperator_DeleteByWorkerId(t *testing.T) {
+	client, err := ConnectMongoDB()
+	if err != nil {
+		log.Fatal(err)
+	}
+	db := NewDBWorkerInstalled(client, database)
+	defer db.client.Disconnect(context.Background())
+	workerid := uuid.NewString()
+	// insert first
+	ws := make([]*WorkerInstalledInfo, 0)
+	for i := 0; i < 10; i++ {
+		w := &WorkerInstalledInfo{
+			WorkerId: workerid,
+			ModelId:  i,
+			GpuFree:  10011,
+		}
+		ws = append(ws, w)
+	}
+	if _, err := db.InsertMany(context.Background(), ws); err != nil {
+		t.Errorf("insert worker failed with err:%s", err)
+	}
+	n, err := db.DeleteByWorkerId(context.Background(), workerid)
+	if err != nil {
+		t.Errorf("delete worker failed with err:%s", err)
+	}
+	if n != len(ws) {
+		t.Errorf("delete worker failed with err:%s", err)
+	}
+
+}
+
+//
+//func TestWorkerInstalledOperator_Get(t *testing.T) {
+//	client, err := ConnectMongoDB()
+//	if err != nil {
+//		log.Fatal(err)
+//	}
+//	db := NewDBWorkerInstalled(client, database)
+//	defer db.client.Disconnect(context.Background())
+//	workerid := uuid.NewString()
+//	// insert first
+//	ws := make([]*WorkerInstalledInfo, 0)
+//	for i := 0; i < 10; i++ {
+//		w := &WorkerInstalledInfo{
+//			WorkerId: workerid,
+//			ModelId:  i,
+//			GpuFree:  10011,
+//		}
+//		ws = append(ws, w)
+//	}
+//	if _, err := db.InsertMany(context.Background(), ws); err != nil {
+//		t.Errorf("insert worker failed with err:%s", err)
+//	}
+//	for _, w := range ws {
+//		worker, err := db.Get(context.Background(), w.ID)
+//		if err != nil {
+//			t.Errorf("get worker failed with err:%s", err)
+//		}
+//		if worker == nil {
+//			t.Errorf("get worker failed with err:%s", err)
+//		}
+//	}
+//
+//}
+
+func TestWorkerInstalledOperator_UpdateGpuFree(t *testing.T) {
+	client, err := ConnectMongoDB()
+	if err != nil {
+		log.Fatal(err)
+	}
+	db := NewDBWorkerInstalled(client, database)
+	defer db.client.Disconnect(context.Background())
+	workerid := uuid.NewString()
+	// insert first
+	ws := make([]*WorkerInstalledInfo, 0)
+	for i := 0; i < 10; i++ {
+		w := &WorkerInstalledInfo{
+			WorkerId: workerid,
+			ModelId:  i,
+			GpuFree:  10011,
+		}
+		ws = append(ws, w)
+	}
+	is, err := db.InsertMany(context.Background(), ws)
+	if err != nil {
+		t.Errorf("insert worker failed with err:%s", err)
+	}
+	ids := make([]string, 0)
+	for _, i := range is.InsertedIDs {
+		ids = append(ids, i.(primitive.ObjectID).Hex())
+	}
+	gpufree := 102222
+	gpuseq := 3
+	if err := db.UpdateGpuFree(context.Background(), workerid, int64(gpufree), gpuseq); err != nil {
+		t.Errorf("update worker failed with err:%s", err)
+	}
+	for _, id := range ids {
+		worker, err := db.Get(context.Background(), id)
 		if err != nil {
-				panic(fmt.Sprintf("insert worker failed with err:%s", err))
+			t.Errorf("get worker failed with err:%s", err)
+		}
+		if worker.GpuFree != int64(gpufree) {
+			t.Errorf("update worker failed with err:%s", err)
 		}
+		if worker.GpuSeq != gpuseq {
+			t.Errorf("update worker failed with err:%s", err)
 		}
 	}
 }
 
+func TestWorkerInstalledOperator_FindWorkerByModelIdAndGpuMem(t *testing.T) {
+	client, err := ConnectMongoDB()
+	if err != nil {
+		log.Fatal(err)
+	}
+	db := NewDBWorkerInstalled(client, database)
+	defer db.client.Disconnect(context.Background())
+	db.Clear()
+
+	// insert first
+	ws := make([]*WorkerInstalledInfo, 0)
+	for i := 0; i < 10; i++ {
+		w := &WorkerInstalledInfo{
+			WorkerId: uuid.NewString(),
+			ModelId:  i,
+			GpuFree:  100,
+		}
+		ws = append(ws, w)
+	}
+	if _, err := db.InsertMany(context.Background(), ws); err != nil {
+		t.Errorf("insert worker failed with err:%s", err)
+	}
+	{
+		id := 1
+		mem := 101
+		workers, err := db.FindWorkerByModelIdAndGpuMem(context.Background(), id, int64(mem), 10)
+		if err != nil {
+			t.Errorf("find worker failed with err:%s", err)
+		} else {
+			if len(workers) != 0 {
+				t.Error("test failed with worker should be nil")
+			}
+		}
+	}
+	{
+		id := 1
+		mem := 90
+		workers, err := db.FindWorkerByModelIdAndGpuMem(context.Background(), id, int64(mem), 10)
+		if err != nil {
+			t.Errorf("find worker failed with err:%s", err)
+		} else {
+			if len(workers) != 1 {
+				t.Error("test failed with worker length should be 1")
+			}
+
+		}
+	}
+}
+
+//
+//func TestWorkerInstalledOperator_Insert(t *testing.T) {
+//	client, err := ConnectMongoDB()
+//	if err != nil {
+//		log.Fatal(err)
+//	}
+//	count := 1
+//	infodb := NewDBWorker(client, database)
+//	defer infodb.client.Disconnect(context.Background())
+//	db := NewDBWorkerInstalled(client, database)
+//	defer db.client.Disconnect(context.Background())
+//	rundb := NewDBWorkerRunning(client, database)
+//	defer rundb.client.Disconnect(context.Background())
+//
+//	for i := 0; i < count; i++ {
+//		worker := generateAWroker()
+//		infodb.InsertWorker(context.Background(), worker)
+//		for _, installed := range worker.Models.InstalledModels {
+//			id, _ := strconv.Atoi(installed.ModelID)
+//			_, err := db.Insert(context.Background(), &WorkerInstalledInfo{
+//				WorkerId: worker.WorkerId,
+//				ModelId:  id,
+//				GpuFree:  10011,
+//			})
+//			if err != nil {
+//				panic(fmt.Sprintf("insert worker failed with err:%s", err))
+//			}
+//		}
+//
+//		for _, running := range worker.Models.RunningModels {
+//			id, _ := strconv.Atoi(running.ModelID)
+//			_, err := rundb.Insert(context.Background(), &WorkerRunningInfo{
+//				WorkerId: worker.WorkerId,
+//				ModelId:  id,
+//				ExecTime: 100,
+//			})
+//			if err != nil {
+//				panic(fmt.Sprintf("insert worker failed with err:%s", err))
+//			}
+//		}
+//	}
+//}
+
 func BenchmarkWorkerInstalledOperator_UpdateGpuFree(b *testing.B) {
 
 	client, err := ConnectMongoDB()
@@ -145,18 +364,3 @@ func BenchmarkWorkerInstalledOperator_FindWorkerByModelIdAndGpuMem_Parallel(b *t
 	})
 
 }
-
-func TestWorkerInstalledOperator_UpdateGpuFree(t *testing.T) {
-	client, err := ConnectMongoDB()
-	if err != nil {
-		log.Fatal(err)
-	}
-	db := NewDBWorkerInstalled(client, database)
-	defer db.client.Disconnect(context.Background())
-	workerid := "ebb6a2df-9e91-4b3c-ace8-748aaeca718c"
-	gpufree := 102222
-	gpuseq := 3
-	if err := db.UpdateGpuFree(context.Background(), workerid, int64(gpufree), gpuseq); err != nil {
-		panic(fmt.Sprintf("update worker failed with err:%s", err))
-	}
-}