add test and fix bug

operator/workerinfo.go

@@ -185,12 +185,8 @@ func (d *WorkerInfoOperator) UpdateGPUUsage(ctx context.Context, workerid string
 	updates := make([]mongo.WriteModel, len(usages))
 	for i, usage := range usages {
 		// update gpu info
-		//MemFree     int64  `bson:"mem_free" json:"mem_free"`
-		//Usage       int    `bson:"usage" json:"usage"`
-		//Temp        int    `bson:"temp" json:"temp "`
-		//PowerRt     int    `bson:"power_rt" json:"power_rt"`
-		updates[i] = mongo.NewUpdateOneModel().SetFilter(bson.M{"worker_id": workerid, "hardware.gpu.seq": usage.GpuSeq}).SetUpdate(bson.M{"$set": bson.M{"hardware.gpu.$.usage": usage.Usage,
-			"hardware.gpu.$.mem_free": usage.MemFree, "hardware.gpu.$.temp": usage.Temp, "hardware.gpu.$.power_rt": usage.Powrt}})
+		updates[i] = mongo.NewUpdateOneModel().SetFilter(bson.M{"worker_id": workerid, "hardware.GPU.seq": usage.GpuSeq}).SetUpdate(bson.M{"$set": bson.M{"hardware.GPU.$.usage": usage.Usage,
+			"hardware.GPU.$.mem_free": usage.MemFree, "hardware.GPU.$.temp": usage.Temp, "hardware.GPU.$.power_rt": usage.Powrt}})
 	}
 	_, err := d.col.BulkWrite(ctx, updates)
 	return err

operator/workerinfo_test.go

@@ -360,7 +360,57 @@ func TestWorkerInfoOperator_AddInstalledModels(t *testing.T) {
 		t.Error("add installed models failed")
 
 	}
+}
 
+func TestWorkerInfoOperator_UpdateGPUUsage(t *testing.T) {
+	client, err := ConnectMongoDB()
+	if err != nil {
+		log.Fatal(err)
+	}
+	db := NewDBWorker(client, database)
+	defer db.client.Disconnect(context.Background())
+	db.Clear()
+	worker := &WorkerInfo{
+		WorkerId: uuid.NewString(),
+		NodeInfo: generateANodeInfo(),
+		Models: &types.ModelInfo{
+			InstalledModels: make([]*types.InstalledModel, 0, 1000),
+			RunningModels: []*types.RunningModel{
+				&types.RunningModel{
+					ModelID:       "1",
+					GpuSeq:        0,
+					GpuRAM:        1024 * 1024 * 1024,
+					StartedTime:   time.Now().Unix(),
+					LastWorkTime:  time.Now().Unix(),
+					TotalRunCount: 100,
+					ExecTime:      100,
+				},
+			},
+		},
+		Hardware: generateAHardware(),
+	}
+	_, err = db.InsertWorker(context.Background(), worker)
+	if err != nil {
+		t.Errorf("insert worker failed with err:%s", err)
+	}
+
+	gpuusages := []types.GpuUsage{
+		{0, 50, 10, 20, 30},
+		{1, 51, 11, 21, 31},
+		{2, 52, 12, 22, 32},
+	}
+
+	// update gpu usage
+	if err := db.UpdateGPUUsage(context.Background(), worker.WorkerId, gpuusages); err != nil {
+		t.Error("update gpu usage failed")
+	}
+	b, err := db.FindWorkerByWorkerId(context.Background(), worker.WorkerId)
+	if err != nil {
+		t.Error("find worker failed")
+	}
+	if b.Hardware.GPU[0].Usage != gpuusages[0].Usage {
+		t.Error("update gpu usage compare failed")
+	}
 }
 
 func TestWorkerInfoOperator_AddRunningModels(t *testing.T) {