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