task_msg_test.go 2.77 KB
package test

import (
	"encoding/json"
	"example.com/m/models"
	"example.com/m/nm"
	"example.com/m/operate"
	"fmt"
	"github.com/golang/groupcache/lru"
	nodeManagerV1 "github.com/odysseus/odysseus-protocol/gen/proto/go/nodemanager/v1"
	"net/http"
	"sync"
	"testing"
)

func TestJson(t *testing.T) {
	type DockerCmd struct {
		ContainerPort string `json:"container_port"`
		HostIp        string
		HostPort      string
	}

	type TaskCmd struct {
		ImageName string     `json:"image_name"`
		DockerCmd *DockerCmd `json:"docker_cmd"`
		ApiUrl    string     `json:"api_url"`
	}

	taskCmd := &TaskCmd{
		ImageName: "onlydd/llm-server:0119",
		DockerCmd: &DockerCmd{
			ContainerPort: "8888",
		},
		ApiUrl: "https://192.168.1.120:8888/llm/test/get/sign",
	}

	marshal, err := json.Marshal(taskCmd)
	if err != nil {
		_ = fmt.Errorf("error marshalling task cmd: %s", err.Error())
		return
	}

	fmt.Println("marshal:", string(marshal))
}

func TestTaskHandler_computeTaskHandler(t1 *testing.T) {
	type fields struct {
		wg             *sync.WaitGroup
		lruCache       *lru.Cache
		DockerOp       *operate.DockerOp
		CmdOp          *operate.Command
		TaskMsg        chan *nodeManagerV1.PushTaskMessage
		TaskRespHeader map[string][]byte
		TaskRespBody   map[string][]byte
		TaskIsSuccess  map[string]bool
		HttpClient     *http.Client
	}

	type args struct {
		taskMsg *nodeManagerV1.PushTaskMessage
	}
	m := &models.TaskCmd{
		ImageName: "llm-server:latest",
		DockerCmd: &models.DockerCmd{
			HostIp:   "0.0.0.0",
			HostPort: "",
		},
		ApiUrl: "http://192.168.1.120:%d/llm/test/get/sign",
	}
	marshal, err := json.Marshal(m)
	if err != nil {
		fmt.Println("Error marshalling:", err)
		return
	}

	taskParam := &models.TaskReq{
		TaskId: "22222",
	}
	taskParamBytes, err := json.Marshal(taskParam)
	if err != nil {
		fmt.Println("Error marshalling:", err)
		return
	}
	n := args{
		taskMsg: &nodeManagerV1.PushTaskMessage{
			Workload:  111,
			TaskCmd:   string(marshal),
			TaskParam: taskParamBytes,
		},
	}

	tests := []struct {
		name   string
		fields fields
		args   args
	}{
		{
			"test send task",
			fields{
				wg:             &sync.WaitGroup{},
				lruCache:       lru.New(100),
				DockerOp:       operate.NewDockerOp(),
				TaskMsg:        make(chan *nodeManagerV1.PushTaskMessage, 0),
				TaskRespHeader: make(map[string][]byte, 0),
				TaskRespBody:   make(map[string][]byte, 0),
				TaskIsSuccess:  make(map[string]bool, 0),
				HttpClient:     &http.Client{},
			},
			n,
		},
	}
	for _, tt := range tests {
		t1.Run(tt.name, func(t1 *testing.T) {
			t := &nm.TaskHandler{
				Wg:         tt.fields.wg,
				LruCache:   tt.fields.lruCache,
				DockerOp:   tt.fields.DockerOp,
				TaskMsg:    tt.fields.TaskMsg,
				HttpClient: tt.fields.HttpClient,
			}
			tt.fields.wg.Add(1)
			t.ComputeTaskHandler(tt.args.taskMsg)
		})
	}
}