1package policy
  2
  3import (
  4	"os"
  5	"path"
  6	"testing"
  7
  8	"postern/internal/db"
  9
 10	"go.starlark.net/starlark"
 11)
 12
 13func TestKVModule(t *testing.T) {
 14	posternDB, err := db.OpenDB(&db.Config{
 15		DBPath: path.Join(t.TempDir(), "meta.db"),
 16	})
 17	if err != nil {
 18		t.Fatal(err)
 19	}
 20
 21	defer func(posternDB *db.DB) {
 22		err := posternDB.Close()
 23		if err != nil {
 24			t.Fatal(err)
 25		}
 26	}(posternDB)
 27
 28	policy := `
 29def init():
 30	return None
 31
 32def test_kv_set(k, v):
 33	kv.set(k, v)
 34
 35def test_kv_get(k):
 36	return kv.get(k)
 37
 38def test_kv_get_default(k, d):
 39	return kv.get(k, default=d)
 40
 41def test_kv_inc_delta(k, d):
 42	return kv.inc(k, delta=d)
 43
 44def test_kv_inc(k):
 45	return kv.inc(k)
 46
 47def test_kv_pop(k):
 48	return kv.pop(k)
 49`
 50
 51	policyFile := path.Join(t.TempDir(), "policy.star")
 52	err = os.WriteFile(policyFile, []byte(policy), 0644)
 53	if err != nil {
 54		t.Fatal(err)
 55	}
 56
 57	e, err := NewEngine(policyFile, posternDB)
 58	if err != nil {
 59		t.Fatal(err)
 60	}
 61
 62	// KV Set initial value
 63	hookVal, ok := e.globals["test_kv_set"]
 64	if !ok {
 65		t.Fatal("test_kv_set not found in engine globals")
 66	}
 67
 68	k, v := starlark.String("key"), starlark.MakeInt(1)
 69	_, err = starlark.Call(e.thread, hookVal, starlark.Tuple{k, v}, nil)
 70	if err != nil {
 71		t.Fatal(err)
 72	}
 73
 74	// KV Set unsupported type
 75	k, unsupported := starlark.String("key"), &starlark.Dict{}
 76	_, err = starlark.Call(e.thread, hookVal, starlark.Tuple{k, unsupported}, nil)
 77	if err == nil {
 78		t.Fatal("expected unsupported type to throw an error")
 79	}
 80
 81	// KV get initial value
 82	hookVal, ok = e.globals["test_kv_get"]
 83	if !ok {
 84		t.Fatal("test_kv_get not found in engine globals")
 85	}
 86
 87	res, err := starlark.Call(e.thread, hookVal, starlark.Tuple{k}, nil)
 88	if err != nil {
 89		t.Fatal(err)
 90	}
 91	if res != v {
 92		t.Fatalf("got %s, want %s", res, v)
 93	}
 94
 95	// Get default value for non-existing key
 96	hookVal, ok = e.globals["test_kv_get_default"]
 97	if !ok {
 98		t.Fatal("test_kv_get_default not found in engine globals")
 99	}
100	res, err = starlark.Call(e.thread, hookVal, starlark.Tuple{starlark.String("not-exist"), starlark.String("default")}, nil)
101	if err != nil {
102		t.Fatal(err)
103	}
104	if want := starlark.String("default"); res != want {
105		t.Fatalf("got %s, want %s", res, want)
106	}
107
108	// KV increment initial value by default delta=1
109	hookVal, ok = e.globals["test_kv_inc"]
110	if !ok {
111		t.Fatal("test_kv_inc not found in engine globals")
112	}
113
114	res, err = starlark.Call(e.thread, hookVal, starlark.Tuple{k}, nil)
115	if err != nil {
116		t.Fatal(err)
117	}
118	if want := starlark.MakeInt(2); res != want {
119		t.Fatalf("got %s, want %s", res, want)
120	}
121
122	// KV increment by delta=2
123	hookVal, ok = e.globals["test_kv_inc_delta"]
124	if !ok {
125		t.Fatal("test_kv_inc not found in engine globals")
126	}
127
128	res, err = starlark.Call(e.thread, hookVal, starlark.Tuple{k, starlark.MakeInt(2)}, nil)
129	if err != nil {
130		t.Fatal(err)
131	}
132	if want := starlark.MakeInt(4); res != want {
133		t.Fatalf("got %s, want %s", res, want)
134	}
135
136	// KV pop
137	hookVal, ok = e.globals["test_kv_pop"]
138	if !ok {
139		t.Fatal("test_kv_pop not found in engine globals")
140	}
141
142	res, err = starlark.Call(e.thread, hookVal, starlark.Tuple{k}, nil)
143	if err != nil {
144		t.Fatal(err)
145	}
146	if want := starlark.MakeInt(4); res != want {
147		t.Fatalf("got %s, want %s", res, want)
148	}
149
150	// KV pop non-existing
151	res, err = starlark.Call(e.thread, hookVal, starlark.Tuple{k}, nil)
152	if err != nil {
153		t.Fatal(err)
154	}
155	if res != starlark.None {
156		t.Fatalf("got %s, want %s", res, starlark.None)
157	}
158}