xgo/script_test.go

667 lines
15 KiB
Go

package tengo_test
import (
"context"
"errors"
"fmt"
"math/rand"
"strconv"
"strings"
"sync"
"testing"
"time"
"github.com/d5/tengo/v2"
"github.com/d5/tengo/v2/require"
"github.com/d5/tengo/v2/stdlib"
"github.com/d5/tengo/v2/token"
)
func TestScript_Add(t *testing.T) {
s := tengo.NewScript([]byte(`a := b; c := test(b); d := test(5)`))
require.NoError(t, s.Add("b", 5)) // b = 5
require.NoError(t, s.Add("b", "foo")) // b = "foo" (re-define before compilation)
require.NoError(t, s.Add("test",
func(args ...tengo.Object) (ret tengo.Object, err error) {
if len(args) > 0 {
switch arg := args[0].(type) {
case *tengo.Int:
return &tengo.Int{Value: arg.Value + 1}, nil
}
}
return &tengo.Int{Value: 0}, nil
}))
c, err := s.Compile()
require.NoError(t, err)
require.NoError(t, c.Run())
require.Equal(t, "foo", c.Get("a").Value())
require.Equal(t, "foo", c.Get("b").Value())
require.Equal(t, int64(0), c.Get("c").Value())
require.Equal(t, int64(6), c.Get("d").Value())
}
func TestScript_Remove(t *testing.T) {
s := tengo.NewScript([]byte(`a := b`))
err := s.Add("b", 5)
require.NoError(t, err)
require.True(t, s.Remove("b")) // b is removed
_, err = s.Compile() // should not compile because b is undefined
require.Error(t, err)
}
func TestScript_Run(t *testing.T) {
s := tengo.NewScript([]byte(`a := b`))
err := s.Add("b", 5)
require.NoError(t, err)
c, err := s.Run()
require.NoError(t, err)
require.NotNil(t, c)
compiledGet(t, c, "a", int64(5))
}
func TestScript_BuiltinModules(t *testing.T) {
s := tengo.NewScript([]byte(`math := import("math"); a := math.abs(-19.84)`))
s.SetImports(stdlib.GetModuleMap("math"))
c, err := s.Run()
require.NoError(t, err)
require.NotNil(t, c)
compiledGet(t, c, "a", 19.84)
c, err = s.Run()
require.NoError(t, err)
require.NotNil(t, c)
compiledGet(t, c, "a", 19.84)
s.SetImports(stdlib.GetModuleMap("os"))
_, err = s.Run()
require.Error(t, err)
s.SetImports(nil)
_, err = s.Run()
require.Error(t, err)
}
func TestScript_SourceModules(t *testing.T) {
s := tengo.NewScript([]byte(`
enum := import("enum")
a := enum.all([1,2,3], func(_, v) {
return v > 0
})
`))
s.SetImports(stdlib.GetModuleMap("enum"))
c, err := s.Run()
require.NoError(t, err)
require.NotNil(t, c)
compiledGet(t, c, "a", true)
s.SetImports(nil)
_, err = s.Run()
require.Error(t, err)
}
func TestScript_SetMaxConstObjects(t *testing.T) {
// one constant '5'
s := tengo.NewScript([]byte(`a := 5`))
s.SetMaxConstObjects(1) // limit = 1
_, err := s.Compile()
require.NoError(t, err)
s.SetMaxConstObjects(0) // limit = 0
_, err = s.Compile()
require.Error(t, err)
require.Equal(t, "exceeding constant objects limit: 1", err.Error())
// two constants '5' and '1'
s = tengo.NewScript([]byte(`a := 5 + 1`))
s.SetMaxConstObjects(2) // limit = 2
_, err = s.Compile()
require.NoError(t, err)
s.SetMaxConstObjects(1) // limit = 1
_, err = s.Compile()
require.Error(t, err)
require.Equal(t, "exceeding constant objects limit: 2", err.Error())
// duplicates will be removed
s = tengo.NewScript([]byte(`a := 5 + 5`))
s.SetMaxConstObjects(1) // limit = 1
_, err = s.Compile()
require.NoError(t, err)
s.SetMaxConstObjects(0) // limit = 0
_, err = s.Compile()
require.Error(t, err)
require.Equal(t, "exceeding constant objects limit: 1", err.Error())
// no limit set
s = tengo.NewScript([]byte(`a := 1 + 2 + 3 + 4 + 5`))
_, err = s.Compile()
require.NoError(t, err)
}
func TestScriptConcurrency(t *testing.T) {
solve := func(a, b, c int) (d, e int) {
a += 2
b += c
a += b * 2
d = a + b + c
e = 0
for i := 1; i <= d; i++ {
e += i
}
e *= 2
return
}
code := []byte(`
mod1 := import("mod1")
a += 2
b += c
a += b * 2
arr := [a, b, c]
arrstr := string(arr)
map := {a: a, b: b, c: c}
d := a + b + c
s := 0
for i:=1; i<=d; i++ {
s += i
}
e := mod1.double(s)
`)
mod1 := map[string]tengo.Object{
"double": &tengo.UserFunction{
Value: func(args ...tengo.Object) (
ret tengo.Object,
err error,
) {
arg0, _ := tengo.ToInt64(args[0])
ret = &tengo.Int{Value: arg0 * 2}
return
},
},
}
scr := tengo.NewScript(code)
_ = scr.Add("a", 0)
_ = scr.Add("b", 0)
_ = scr.Add("c", 0)
mods := tengo.NewModuleMap()
mods.AddBuiltinModule("mod1", mod1)
scr.SetImports(mods)
compiled, err := scr.Compile()
require.NoError(t, err)
executeFn := func(compiled *tengo.Compiled, a, b, c int) (d, e int) {
_ = compiled.Set("a", a)
_ = compiled.Set("b", b)
_ = compiled.Set("c", c)
err := compiled.Run()
require.NoError(t, err)
d = compiled.Get("d").Int()
e = compiled.Get("e").Int()
return
}
concurrency := 500
var wg sync.WaitGroup
wg.Add(concurrency)
for i := 0; i < concurrency; i++ {
go func(compiled *tengo.Compiled) {
time.Sleep(time.Duration(rand.Int63n(50)) * time.Millisecond)
defer wg.Done()
a := rand.Intn(10)
b := rand.Intn(10)
c := rand.Intn(10)
d, e := executeFn(compiled, a, b, c)
expectedD, expectedE := solve(a, b, c)
require.Equal(t, expectedD, d, "input: %d, %d, %d", a, b, c)
require.Equal(t, expectedE, e, "input: %d, %d, %d", a, b, c)
}(compiled.Clone())
}
wg.Wait()
}
type Counter struct {
tengo.ObjectImpl
value int64
}
func (o *Counter) TypeName() string {
return "counter"
}
func (o *Counter) String() string {
return fmt.Sprintf("Counter(%d)", o.value)
}
func (o *Counter) BinaryOp(
op token.Token,
rhs tengo.Object,
) (tengo.Object, error) {
switch rhs := rhs.(type) {
case *Counter:
switch op {
case token.Add:
return &Counter{value: o.value + rhs.value}, nil
case token.Sub:
return &Counter{value: o.value - rhs.value}, nil
}
case *tengo.Int:
switch op {
case token.Add:
return &Counter{value: o.value + rhs.Value}, nil
case token.Sub:
return &Counter{value: o.value - rhs.Value}, nil
}
}
return nil, errors.New("invalid operator")
}
func (o *Counter) IsFalsy() bool {
return o.value == 0
}
func (o *Counter) Equals(t tengo.Object) bool {
if tc, ok := t.(*Counter); ok {
return o.value == tc.value
}
return false
}
func (o *Counter) Copy() tengo.Object {
return &Counter{value: o.value}
}
func (o *Counter) Call(_ ...tengo.Object) (tengo.Object, error) {
return &tengo.Int{Value: o.value}, nil
}
func (o *Counter) CanCall() bool {
return true
}
func TestScript_CustomObjects(t *testing.T) {
c := compile(t, `a := c1(); s := string(c1); c2 := c1; c2++`, M{
"c1": &Counter{value: 5},
})
compiledRun(t, c)
compiledGet(t, c, "a", int64(5))
compiledGet(t, c, "s", "Counter(5)")
compiledGetCounter(t, c, "c2", &Counter{value: 6})
c = compile(t, `
arr := [1, 2, 3, 4]
for x in arr {
c1 += x
}
out := c1()
`, M{
"c1": &Counter{value: 5},
})
compiledRun(t, c)
compiledGet(t, c, "out", int64(15))
}
func compiledGetCounter(
t *testing.T,
c *tengo.Compiled,
name string,
expected *Counter,
) {
v := c.Get(name)
require.NotNil(t, v)
actual := v.Value().(*Counter)
require.NotNil(t, actual)
require.Equal(t, expected.value, actual.value)
}
func TestScriptSourceModule(t *testing.T) {
// script1 imports "mod1"
scr := tengo.NewScript([]byte(`out := import("mod")`))
mods := tengo.NewModuleMap()
mods.AddSourceModule("mod", []byte(`export 5`))
scr.SetImports(mods)
c, err := scr.Run()
require.NoError(t, err)
require.Equal(t, int64(5), c.Get("out").Value())
// executing module function
scr = tengo.NewScript([]byte(`fn := import("mod"); out := fn()`))
mods = tengo.NewModuleMap()
mods.AddSourceModule("mod",
[]byte(`a := 3; export func() { return a + 5 }`))
scr.SetImports(mods)
c, err = scr.Run()
require.NoError(t, err)
require.Equal(t, int64(8), c.Get("out").Value())
scr = tengo.NewScript([]byte(`out := import("mod")`))
mods = tengo.NewModuleMap()
mods.AddSourceModule("mod",
[]byte(`text := import("text"); export text.title("foo")`))
mods.AddBuiltinModule("text",
map[string]tengo.Object{
"title": &tengo.UserFunction{
Name: "title",
Value: func(args ...tengo.Object) (tengo.Object, error) {
s, _ := tengo.ToString(args[0])
return &tengo.String{Value: strings.Title(s)}, nil
}},
})
scr.SetImports(mods)
c, err = scr.Run()
require.NoError(t, err)
require.Equal(t, "Foo", c.Get("out").Value())
scr.SetImports(nil)
_, err = scr.Run()
require.Error(t, err)
}
func BenchmarkArrayIndex(b *testing.B) {
bench(b.N, `a := [1, 2, 3, 4, 5, 6, 7, 8, 9];
for i := 0; i < 1000; i++ {
a[0]; a[1]; a[2]; a[3]; a[4]; a[5]; a[6]; a[7]; a[7];
}
`)
}
func BenchmarkArrayIndexCompare(b *testing.B) {
bench(b.N, `a := [1, 2, 3, 4, 5, 6, 7, 8, 9];
for i := 0; i < 1000; i++ {
1; 2; 3; 4; 5; 6; 7; 8; 9;
}
`)
}
func bench(n int, input string) {
s := tengo.NewScript([]byte(input))
c, err := s.Compile()
if err != nil {
panic(err)
}
for i := 0; i < n; i++ {
if err := c.Run(); err != nil {
panic(err)
}
}
}
type M map[string]interface{}
func TestCompiled_Get(t *testing.T) {
// simple script
c := compile(t, `a := 5`, nil)
compiledRun(t, c)
compiledGet(t, c, "a", int64(5))
// user-defined variables
compileError(t, `a := b`, nil) // compile error because "b" is not defined
c = compile(t, `a := b`, M{"b": "foo"}) // now compile with b = "foo" defined
compiledGet(t, c, "a", nil) // a = undefined; because it's before Compiled.Run()
compiledRun(t, c) // Compiled.Run()
compiledGet(t, c, "a", "foo") // a = "foo"
}
func TestCompiled_GetAll(t *testing.T) {
c := compile(t, `a := 5`, nil)
compiledRun(t, c)
compiledGetAll(t, c, M{"a": int64(5)})
c = compile(t, `a := b`, M{"b": "foo"})
compiledRun(t, c)
compiledGetAll(t, c, M{"a": "foo", "b": "foo"})
c = compile(t, `a := b; b = 5`, M{"b": "foo"})
compiledRun(t, c)
compiledGetAll(t, c, M{"a": "foo", "b": int64(5)})
}
func TestCompiled_IsDefined(t *testing.T) {
c := compile(t, `a := 5`, nil)
compiledIsDefined(t, c, "a", false) // a is not defined before Run()
compiledRun(t, c)
compiledIsDefined(t, c, "a", true)
compiledIsDefined(t, c, "b", false)
}
func TestCompiled_Set(t *testing.T) {
c := compile(t, `a := b`, M{"b": "foo"})
compiledRun(t, c)
compiledGet(t, c, "a", "foo")
// replace value of 'b'
err := c.Set("b", "bar")
require.NoError(t, err)
compiledRun(t, c)
compiledGet(t, c, "a", "bar")
// try to replace undefined variable
err = c.Set("c", 1984)
require.Error(t, err) // 'c' is not defined
// case #2
c = compile(t, `
a := func() {
return func() {
return b + 5
}()
}()`, M{"b": 5})
compiledRun(t, c)
compiledGet(t, c, "a", int64(10))
err = c.Set("b", 10)
require.NoError(t, err)
compiledRun(t, c)
compiledGet(t, c, "a", int64(15))
}
func TestCompiled_RunContext(t *testing.T) {
// machine completes normally
c := compile(t, `a := 5`, nil)
err := c.RunContext(context.Background())
require.NoError(t, err)
compiledGet(t, c, "a", int64(5))
// timeout
c = compile(t, `for true {}`, nil)
ctx, cancel := context.WithTimeout(context.Background(),
1*time.Millisecond)
defer cancel()
err = c.RunContext(ctx)
require.Equal(t, context.DeadlineExceeded, err)
}
func TestCompiled_CustomObject(t *testing.T) {
c := compile(t, `r := (t<130)`, M{"t": &customNumber{value: 123}})
compiledRun(t, c)
compiledGet(t, c, "r", true)
c = compile(t, `r := (t>13)`, M{"t": &customNumber{value: 123}})
compiledRun(t, c)
compiledGet(t, c, "r", true)
}
// customNumber is a user defined object that can compare to tengo.Int
// very shitty implementation, just to test that token.Less and token.Greater in BinaryOp works
type customNumber struct {
tengo.ObjectImpl
value int64
}
func (n *customNumber) TypeName() string {
return "Number"
}
func (n *customNumber) String() string {
return strconv.FormatInt(n.value, 10)
}
func (n *customNumber) BinaryOp(op token.Token, rhs tengo.Object) (tengo.Object, error) {
tengoInt, ok := rhs.(*tengo.Int)
if !ok {
return nil, tengo.ErrInvalidOperator
}
return n.binaryOpInt(op, tengoInt)
}
func (n *customNumber) binaryOpInt(op token.Token, rhs *tengo.Int) (tengo.Object, error) {
i := n.value
switch op {
case token.Less:
if i < rhs.Value {
return tengo.TrueValue, nil
}
return tengo.FalseValue, nil
case token.Greater:
if i > rhs.Value {
return tengo.TrueValue, nil
}
return tengo.FalseValue, nil
case token.LessEq:
if i <= rhs.Value {
return tengo.TrueValue, nil
}
return tengo.FalseValue, nil
case token.GreaterEq:
if i >= rhs.Value {
return tengo.TrueValue, nil
}
return tengo.FalseValue, nil
}
return nil, tengo.ErrInvalidOperator
}
func TestScript_ImportError(t *testing.T) {
m := `
exp := import("expression")
r := exp(ctx)
`
src := `
export func(ctx) {
closure := func() {
if ctx.actiontimes < 0 { // an error is thrown here because actiontimes is undefined
return true
}
return false
}
return closure()
}`
s := tengo.NewScript([]byte(m))
mods := tengo.NewModuleMap()
mods.AddSourceModule("expression", []byte(src))
s.SetImports(mods)
err := s.Add("ctx", map[string]interface{}{
"ctx": 12,
})
require.NoError(t, err)
_, err = s.Run()
require.True(t, strings.Contains(err.Error(), "expression:4:6"))
}
func compile(t *testing.T, input string, vars M) *tengo.Compiled {
s := tengo.NewScript([]byte(input))
for vn, vv := range vars {
err := s.Add(vn, vv)
require.NoError(t, err)
}
c, err := s.Compile()
require.NoError(t, err)
require.NotNil(t, c)
return c
}
func compileError(t *testing.T, input string, vars M) {
s := tengo.NewScript([]byte(input))
for vn, vv := range vars {
err := s.Add(vn, vv)
require.NoError(t, err)
}
_, err := s.Compile()
require.Error(t, err)
}
func compiledRun(t *testing.T, c *tengo.Compiled) {
err := c.Run()
require.NoError(t, err)
}
func compiledGet(
t *testing.T,
c *tengo.Compiled,
name string,
expected interface{},
) {
v := c.Get(name)
require.NotNil(t, v)
require.Equal(t, expected, v.Value())
}
func compiledGetAll(
t *testing.T,
c *tengo.Compiled,
expected M,
) {
vars := c.GetAll()
require.Equal(t, len(expected), len(vars))
for k, v := range expected {
var found bool
for _, e := range vars {
if e.Name() == k {
require.Equal(t, v, e.Value())
found = true
}
}
require.True(t, found, "variable '%s' not found", k)
}
}
func compiledIsDefined(
t *testing.T,
c *tengo.Compiled,
name string,
expected bool,
) {
require.Equal(t, expected, c.IsDefined(name))
}
func TestCompiled_Clone(t *testing.T) {
script := tengo.NewScript([]byte(`
count += 1
data["b"] = 2
`))
err := script.Add("data", map[string]interface{}{"a": 1})
require.NoError(t, err)
err = script.Add("count", 1000)
require.NoError(t, err)
compiled, err := script.Compile()
require.NoError(t, err)
clone := compiled.Clone()
err = clone.RunContext(context.Background())
require.NoError(t, err)
require.Equal(t, 1000, compiled.Get("count").Int())
require.Equal(t, 1, len(compiled.Get("data").Map()))
require.Equal(t, 1001, clone.Get("count").Int())
require.Equal(t, 2, len(clone.Get("data").Map()))
}