api/mgragent/task_handler_test.go - master - zhendosina/obagent

api/mgragent/task_handler_test.go
105 строк3 KB
zhj-luo
Fix file header
14 авг 2023, 16:10
14 авг 2023, 16:108b4d999
/*
 * Copyright (c) 2023 OceanBase
 * OBAgent is licensed under Mulan PSL v2.
 * You can use this software according to the terms and conditions of the Mulan PSL v2.
 * You may obtain a copy of Mulan PSL v2 at:
 *          http://license.coscl.org.cn/MulanPSL2
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
 * EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
 * MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
 * See the Mulan PSL v2 for more details.
 */

package mgragent

import (
	"context"
	"net/http"
	"net/http/httptest"
	"os"
	"strings"
	"testing"

	"github.com/gin-gonic/gin"

	"github.com/oceanbase/obagent/api/common"
	"github.com/oceanbase/obagent/executor/agent"
	"github.com/oceanbase/obagent/lib/command"
	http2 "github.com/oceanbase/obagent/lib/http"
	path2 "github.com/oceanbase/obagent/lib/path"
)

type S struct {
	agent.TaskToken
	A string
}

func TestAsyncCommandHandler(t *testing.T) {
	os.MkdirAll(path2.TaskStoreDir(), 0755)
	defer os.RemoveAll(path2.TaskStoreDir())

	h := asyncCommandHandler(command.WrapFunc(func(s S) S {
		s.A = s.A + s.A
		return s
	}))
	req, _ := http.NewRequest("POST", "/xxx", strings.NewReader(`{"A":"a", "taskToken":"token12345"}`))
	rec := httptest.NewRecorder()
	ctx, _ := gin.CreateTestContext(rec)
	ctx.Request = req
	ctx.Keys = map[string]interface{}{common.TraceIdKey: "a"}
	h(ctx)
	resp := ctx.Keys[common.OcpAgentResponseKey].(http2.OcpAgentResponse)
	if !resp.Successful || resp.Status != 200 {
		t.Errorf("Fail %+v", resp)
		return
	}
	tokenResult := resp.Data.(TaskTokenResult)
	if tokenResult.TaskToken != "token12345" {
		t.Errorf("bad result %+v", tokenResult)
		return
	}
	result, ok := taskExecutor.WaitResult(command.ExecutionTokenFromString("token12345"))
	if !ok {
		t.Error("wait result failed")
		return
	}
	s := result.Result.(S)
	if s.A != "aa" {
		t.Errorf("bad result %+v", s)
	}
}

func TestAsyncCommandHandler2(t *testing.T) {
	os.MkdirAll(path2.TaskStoreDir(), 0755)
	defer os.RemoveAll(path2.TaskStoreDir())

	h := asyncCommandHandler(command.WrapFunc(func(ctx context.Context, s S) (S, error) {
		s.A = s.A + s.A
		return s, nil
	}))
	req, _ := http.NewRequest("POST", "/xxx", strings.NewReader(`{"A":"a", "taskToken":"token12345"}`))
	rec := httptest.NewRecorder()
	ctx, _ := gin.CreateTestContext(rec)
	ctx.Request = req
	ctx.Keys = map[string]interface{}{common.TraceIdKey: "a"}
	h(ctx)
	resp := ctx.Keys[common.OcpAgentResponseKey].(http2.OcpAgentResponse)
	if !resp.Successful || resp.Status != 200 {
		t.Errorf("Fail %+v", resp)
		return
	}
	tokenResult := resp.Data.(TaskTokenResult)
	if tokenResult.TaskToken != "token12345" {
		t.Errorf("bad result %+v", tokenResult)
		return
	}
	result, ok := taskExecutor.WaitResult(command.ExecutionTokenFromString("token12345"))
	if !ok {
		t.Error("wait result failed")
		return
	}
	s := result.Result.(S)
	if s.A != "aa" {
		t.Errorf("bad result %+v", s)
	}
}