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	s := e.state.Load()
 62
 63	// KV Set initial value
 64	hookVal, ok := s.globals["test_kv_set"]
 65	if !ok {
 66		t.Fatal("test_kv_set not found in engine globals")
 67	}
 68
 69	k, v := starlark.String("key"), starlark.MakeInt(1)
 70	_, err = starlark.Call(s.thread, hookVal, starlark.Tuple{k, v}, nil)
 71	if err != nil {
 72		t.Fatal(err)
 73	}
 74
 75	// KV Set unsupported type
 76	k, unsupported := starlark.String("key"), &starlark.Dict{}
 77	_, err = starlark.Call(s.thread, hookVal, starlark.Tuple{k, unsupported}, nil)
 78	if err == nil {
 79		t.Fatal("expected unsupported type to throw an error")
 80	}
 81
 82	// KV get initial value
 83	hookVal, ok = s.globals["test_kv_get"]
 84	if !ok {
 85		t.Fatal("test_kv_get not found in engine globals")
 86	}
 87
 88	res, err := starlark.Call(s.thread, hookVal, starlark.Tuple{k}, nil)
 89	if err != nil {
 90		t.Fatal(err)
 91	}
 92	if res != v {
 93		t.Fatalf("got %s, want %s", res, v)
 94	}
 95
 96	// Get default value for non-existing key
 97	hookVal, ok = s.globals["test_kv_get_default"]
 98	if !ok {
 99		t.Fatal("test_kv_get_default not found in engine globals")
100	}
101	res, err = starlark.Call(s.thread, hookVal, starlark.Tuple{starlark.String("not-exist"), starlark.String("default")}, nil)
102	if err != nil {
103		t.Fatal(err)
104	}
105	if want := starlark.String("default"); res != want {
106		t.Fatalf("got %s, want %s", res, want)
107	}
108
109	// KV increment initial value by default delta=1
110	hookVal, ok = s.globals["test_kv_inc"]
111	if !ok {
112		t.Fatal("test_kv_inc not found in engine globals")
113	}
114
115	res, err = starlark.Call(s.thread, hookVal, starlark.Tuple{k}, nil)
116	if err != nil {
117		t.Fatal(err)
118	}
119	if want := starlark.MakeInt(2); res != want {
120		t.Fatalf("got %s, want %s", res, want)
121	}
122
123	// KV increment by delta=2
124	hookVal, ok = s.globals["test_kv_inc_delta"]
125	if !ok {
126		t.Fatal("test_kv_inc not found in engine globals")
127	}
128
129	res, err = starlark.Call(s.thread, hookVal, starlark.Tuple{k, starlark.MakeInt(2)}, nil)
130	if err != nil {
131		t.Fatal(err)
132	}
133	if want := starlark.MakeInt(4); res != want {
134		t.Fatalf("got %s, want %s", res, want)
135	}
136
137	// KV pop
138	hookVal, ok = s.globals["test_kv_pop"]
139	if !ok {
140		t.Fatal("test_kv_pop not found in engine globals")
141	}
142
143	res, err = starlark.Call(s.thread, hookVal, starlark.Tuple{k}, nil)
144	if err != nil {
145		t.Fatal(err)
146	}
147	if want := starlark.MakeInt(4); res != want {
148		t.Fatalf("got %s, want %s", res, want)
149	}
150
151	// KV pop non-existing
152	res, err = starlark.Call(s.thread, hookVal, starlark.Tuple{k}, nil)
153	if err != nil {
154		t.Fatal(err)
155	}
156	if res != starlark.None {
157		t.Fatalf("got %s, want %s", res, starlark.None)
158	}
159}