feat: add mocks for testing the llm function (#998)

1. Add `mock.NewArgumentsContext()` for testing llm function.
2. Add `WriteRecord.LLMResult` for retrieving llm function written data.

cli/template/go/init_llm_test.tmpl

@@ -1,36 +1,32 @@
 package main
 
 import (
-	"fmt"
 	"reflect"
 	"testing"
 
-	"github.com/yomorun/yomo/ai"
 	"github.com/yomorun/yomo/serverless/mock"
 )
 
 func TestHandler(t *testing.T) {
 	tests := []struct {
 		name string
 		ctx  *mock.MockContext
-		// want is the expected data and tag that be written by ctx.Write
-		want []mock.WriteRecord
+		// want is the expected result written by ctx.WriteLLMResult()
+		want string
 	}{
 		{
 			name: "get weather",
-			ctx:  mock.NewMockContext([]byte(`{"arguments":"{\"city\":\"New York\",\"latitude\":40.7128,\"longitude\":-74.0060}"}`), 0x33),
-			want: []mock.WriteRecord{
-				{Data: []byte(`{"result":"The current weather in New York (40.712800,-74.006000) is sunny","arguments":"{\"city\":\"New York\",\"latitude\":40.7128,\"longitude\":-74.0060}","is_ok":true}`), Tag: ai.ReducerTag},
-			},
+			ctx:  mock.NewArgumentsContext(`{"city":"New York","latitude":40.7128,"longitude":-74.0060}`, 0x33),
+			want: "The current weather in New York (40.712800,-74.006000) is sunny",
 		},
 		// TODO: add more test cases.
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			Handler(tt.ctx)
-			got := tt.ctx.RecordsWritten()
 
-			fmt.Println(string(got[0].Data))
+			records := tt.ctx.RecordsWritten()
+			got := records[0].LLMResult
 
 			if !reflect.DeepEqual(got, tt.want) {
 				t.Errorf("TestHandler got: %v, want: %v", got, tt.want)

cli/template/template.go

@@ -25,16 +25,40 @@ var (
 
 // get template content
 func GetContent(command string, sfnType string, lang string, isTest bool) ([]byte, error) {
+	name, err := getTemplateFileName(command, sfnType, lang, isTest)
+	if err != nil {
+		return nil, err
+	}
+	f, err := fs.Open(name)
+	if err != nil {
+		if os.IsNotExist(err) {
+			if isTest {
+				return nil, ErrUnsupportedTest
+			}
+			return nil, err
+		}
+		return nil, err
+	}
+	defer f.Close()
+	_, err = f.Stat()
+	if err != nil {
+		return nil, err
+	}
+
+	return fs.ReadFile(name)
+}
+
+func getTemplateFileName(command string, sfnType string, lang string, isTest bool) (string, error) {
 	if command == "" {
 		command = "init"
 	}
 	sfnType, err := validateSfnType(sfnType)
 	if err != nil {
-		return nil, err
+		return "", err
 	}
 	lang, err = validateLang(lang)
 	if err != nil {
-		return nil, err
+		return "", err
 	}
 	sb := new(strings.Builder)
 	sb.WriteString(lang)
@@ -47,25 +71,10 @@ func GetContent(command string, sfnType string, lang string, isTest bool) ([]byt
 	}
 	sb.WriteString(".tmpl")
 
-	// valdiate the path exists
+	// validate the path exists
 	name := sb.String()
-	f, err := fs.Open(name)
-	if err != nil {
-		if os.IsNotExist(err) {
-			if isTest {
-				return nil, ErrUnsupportedTest
-			}
-			return nil, ErrUnsupportedFeature
-		}
-		return nil, err
-	}
-	defer f.Close()
-	_, err = f.Stat()
-	if err != nil {
-		return nil, err
-	}
 
-	return fs.ReadFile(name)
+	return name, nil
 }
 
 func validateSfnType(sfnType string) (string, error) {

cli/template/template_test.go

@@ -0,0 +1,121 @@
+package template
+
+import (
+	"testing"
+)
+
+func TestGetTemplateFileName(t *testing.T) {
+	type args struct {
+		command string
+		sfnType string
+		lang    string
+		isTest  bool
+	}
+	tests := []struct {
+		name    string
+		args    args
+		want    string
+		wantErr bool
+	}{
+		{
+			name: "init_llm_go",
+			args: args{
+				command: "init",
+				sfnType: "llm",
+				lang:    "go",
+				isTest:  false,
+			},
+			want:    "go/init_llm.tmpl",
+			wantErr: false,
+		},
+		{
+			name: "init_normal_node_test",
+			args: args{
+				command: "init",
+				sfnType: "normal",
+				lang:    "node",
+				isTest:  true,
+			},
+			want:    "node/init_normal_test.tmpl",
+			wantErr: false,
+		},
+		{
+			name: "default_command_llm_go",
+			args: args{
+				command: "",
+				sfnType: "llm",
+				lang:    "go",
+				isTest:  false,
+			},
+			want:    "go/init_llm.tmpl",
+			wantErr: false,
+		},
+		{
+			name: "unsupported_sfnType",
+			args: args{
+				command: "init",
+				sfnType: "unsupported",
+				lang:    "go",
+				isTest:  false,
+			},
+			want:    "",
+			wantErr: true,
+		},
+		{
+			name: "unsupported_lang",
+			args: args{
+				command: "init",
+				sfnType: "llm",
+				lang:    "unsupported",
+				isTest:  false,
+			},
+			want:    "",
+			wantErr: true,
+		},
+		{
+			name: "default_sfnType",
+			args: args{
+				command: "init",
+				sfnType: "",
+				lang:    "go",
+				isTest:  false,
+			},
+			want:    "go/init_llm.tmpl",
+			wantErr: false,
+		},
+		{
+			name: "default_lang",
+			args: args{
+				command: "init",
+				sfnType: "llm",
+				lang:    "",
+				isTest:  false,
+			},
+			want:    "go/init_llm.tmpl",
+			wantErr: false,
+		},
+		{
+			name: "default_sfnType_and_lang",
+			args: args{
+				command: "init",
+				sfnType: "",
+				lang:    "",
+				isTest:  false,
+			},
+			want:    "go/init_llm.tmpl",
+			wantErr: false,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got, err := getTemplateFileName(tt.args.command, tt.args.sfnType, tt.args.lang, tt.args.isTest)
+			if (err != nil) != tt.wantErr {
+				t.Errorf("getTemplateFileName() error = %v, wantErr %v", err, tt.wantErr)
+				return
+			}
+			if got != tt.want {
+				t.Errorf("getTemplateFileName() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}

serverless/mock/mock_context.go

@@ -7,16 +7,18 @@ import (
 	"sync"
 
 	"github.com/yomorun/yomo/ai"
+	"github.com/yomorun/yomo/pkg/id"
 	"github.com/yomorun/yomo/serverless"
 )
 
 var _ serverless.Context = (*MockContext)(nil)
 
 // WriteRecord composes the data, tag and target.
 type WriteRecord struct {
-	Data   []byte
-	Tag    uint32
-	Target string
+	Data      []byte
+	Tag       uint32
+	Target    string
+	LLMResult string
 }
 
 // MockContext mock context.
@@ -38,6 +40,23 @@ func NewMockContext(data []byte, tag uint32) *MockContext {
 	}
 }
 
+// NewArgumentsContext creates a Context with the provided arguments and tag.
+// This function is used for testing the LLM function.
+func NewArgumentsContext(arguments string, tag uint32) *MockContext {
+	fnCall := &ai.FunctionCall{
+		Arguments:  arguments,
+		ReqID:      id.New(16),
+		ToolCallID: "chatcmpl-" + id.New(29),
+	}
+	data, _ := fnCall.Bytes()
+
+	return &MockContext{
+		data:   data,
+		tag:    tag,
+		fnCall: fnCall,
+	}
+}
+
 // Data incoming data.
 func (c *MockContext) Data() []byte {
 	return c.data
@@ -123,9 +142,11 @@ func (c *MockContext) WriteLLMResult(result string) error {
 	}
 
 	c.wrSlice = append(c.wrSlice, WriteRecord{
-		Data: buf,
-		Tag:  ai.ReducerTag,
+		Data:      buf,
+		Tag:       ai.ReducerTag,
+		LLMResult: result,
 	})
+
 	return nil
 }