Skip to content

M
mogo
Commit 21dadcd1 authored by vicotor's avatar vicotor
Browse files
Options

update test code

parent 9de47ea9
Hide whitespace changes
Inline Side-by-side
Showing with 262 additions and 45 deletions
operator/workerInstalledInfo.go
View file @ 21dadcd1
...@@ -3,6 +3,7 @@ package operator ...@@ -3,6 +3,7 @@ package operator
import ( import (
"context" "context"
"go.mongodb.org/mongo-driver/bson" "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"
"go.mongodb.org/mongo-driver/mongo/options" "go.mongodb.org/mongo-driver/mongo/options"
) )
...@@ -55,6 +56,18 @@ func (d *WorkerInstalledOperator) DeleteByWorkerId(ctx context.Context, workerid ...@@ -55,6 +56,18 @@ func (d *WorkerInstalledOperator) DeleteByWorkerId(ctx context.Context, workerid
return int(res.DeletedCount), err 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) { //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 // // find all worker that at least one installed model's mode_id is equal modelId
......
operator/workerInstalledInfo_test.go
View file @ 21dadcd1
...@@ -3,54 +3,273 @@ package operator ...@@ -3,54 +3,273 @@ package operator
import ( import (
"context" "context"
"fmt" "fmt"
"github.com/google/uuid"
"go.mongodb.org/mongo-driver/bson/primitive"
"log" "log"
"math/rand" "math/rand"
"strconv"
"testing" "testing"
) )
func TestWorkerInstalledOperator_Insert(t *testing.T) { func TestWorkerInstalledOperator_InsertMany(t *testing.T) {
client, err := ConnectMongoDB() client, err := ConnectMongoDB()
if err != nil { if err != nil {
log.Fatal(err) log.Fatal(err)
} }
count := 1
infodb := NewDBWorker(client, database)
defer infodb.client.Disconnect(context.Background())
db := NewDBWorkerInstalled(client, database) db := NewDBWorkerInstalled(client, database)
defer db.client.Disconnect(context.Background()) defer db.client.Disconnect(context.Background())
rundb := NewDBWorkerRunning(client, database)
defer rundb.client.Disconnect(context.Background()) ws := make([]*WorkerInstalledInfo, 0)
for i := 0; i < 10; i++ {
for i := 0; i < count; i++ { w := &WorkerInstalledInfo{
worker := generateAWroker() WorkerId: uuid.NewString(),
infodb.InsertWorker(context.Background(), worker) ModelId: i,
for _, installed := range worker.Models.InstalledModels { GpuFree: 10011,
id, _ := strconv.Atoi(installed.ModelID) }
_, err := db.Insert(context.Background(), &WorkerInstalledInfo{ ws = append(ws, w)
WorkerId: worker.WorkerId, }
ModelId: id, if res, err := db.InsertMany(context.Background(), ws); err != nil {
GpuFree: 10011, t.Errorf("insert worker failed with err:%s", err)
}) } else {
if err != nil { if len(res.InsertedIDs) != len(ws) {
panic(fmt.Sprintf("insert worker failed with err:%s", err)) t.Errorf("insert worker failed with err:%s", err)
}
} }
}
}
for _, running := range worker.Models.RunningModels { func TestWorkerInstalledOperator_DeleteMany(t *testing.T) {
id, _ := strconv.Atoi(running.ModelID) client, err := ConnectMongoDB()
_, err := rundb.Insert(context.Background(), &WorkerRunningInfo{ if err != nil {
WorkerId: worker.WorkerId, log.Fatal(err)
ModelId: id, }
ExecTime: 100, db := NewDBWorkerInstalled(client, database)
}) defer db.client.Disconnect(context.Background())
if err != nil {
panic(fmt.Sprintf("insert worker failed with err:%s", err)) 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 {
t.Errorf("delete worker failed with err:%s", err)
}
if n != len(models) {
t.Errorf("delete worker failed with err:%s", err)
}
}
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 {
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) { func BenchmarkWorkerInstalledOperator_UpdateGpuFree(b *testing.B) {
client, err := ConnectMongoDB() client, err := ConnectMongoDB()
...@@ -145,18 +364,3 @@ func BenchmarkWorkerInstalledOperator_FindWorkerByModelIdAndGpuMem_Parallel(b *t ...@@ -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))
}
}
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