more tests

Author: fredmaggiowski

audit/agent_log.go

@@ -33,7 +33,8 @@ func NewLogAgent(l logging.Logger) Agent {
 }
 
 func (a *logAgent) Trace(ctx context.Context, auditData Audit) {
-	data := a.cache.Load("")
+	data := a.cache.Load()
+
 	auditData.generateID()
 	if data != nil {
 		auditData.applyDataFromPolicy(data)

audit/cache.go

@@ -17,18 +17,18 @@ package audit
 type Data map[string]interface{}
 
 type AuditCache interface {
-	Store(id string, d Data)
-	Load(id string) Data
+	Store(d Data)
+	Load() Data
 }
 
 type SingleRecordCache struct {
 	data Data
 }
 
-func (c *SingleRecordCache) Store(id string, d Data) {
+func (c *SingleRecordCache) Store(d Data) {
 	c.data = d
 }
 
-func (c *SingleRecordCache) Load(id string) Data {
+func (c *SingleRecordCache) Load() Data {
 	return c.data
 }

audit/cache_test.go

@@ -0,0 +1,33 @@
+// Copyright 2024 Mia srl
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package audit
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestSingleRecordCache(t *testing.T) {
+	t.Run("store + load", func(t *testing.T) {
+		c := SingleRecordCache{}
+		storedData := Data{"key": "val"}
+
+		c.Store(storedData)
+		loadData := c.Load()
+
+		require.Equal(t, storedData, loadData)
+	})
+}

audit/context.go

@@ -28,7 +28,7 @@ func WithAuditCache(ctx context.Context, agent Agent) context.Context {
 func GetAuditCache(ctx context.Context) (AuditCache, error) {
 	auditCache, ok := ctx.Value(cacheKey{}).(AuditCache)
 	if !ok {
-		return nil, fmt.Errorf("todo")
+		return nil, fmt.Errorf("failed to extract audit cache from context")
 	}
 	return auditCache, nil
 }

audit/context_test.go

@@ -0,0 +1,35 @@
+// Copyright 2024 Mia srl
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package audit
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestContext(t *testing.T) {
+	t.Run("add to context and get", func(t *testing.T) {
+		ctx := context.Background()
+
+		agent := &noopAgent{}
+
+		retrievedAgent, err := GetAuditCache(WithAuditCache(ctx, agent))
+
+		require.NoError(t, err)
+		require.NotNil(t, retrievedAgent)
+	})
+}

audit/rego_builtin.go

@@ -15,8 +15,6 @@
 package audit
 
 import (
-	"fmt"
-
 	"github.com/open-policy-agent/opa/ast"
 	"github.com/open-policy-agent/opa/rego"
 	"github.com/open-policy-agent/opa/types"
@@ -26,34 +24,30 @@ var SetLabelsDecl = &ast.Builtin{
 	Name: "set_audit_labels",
 	Decl: types.NewFunction(
 		types.Args(
-			types.S, // id
 			types.A, // payload
 		),
 		types.A,
 	),
 }
 
-var SetLabels = rego.Function2(
+var SetLabels = rego.Function1(
 	&rego.Function{
 		Name: SetLabelsDecl.Name,
 		Decl: SetLabelsDecl.Decl,
 	},
-	func(bctx rego.BuiltinContext, op1, op2 *ast.Term) (*ast.Term, error) {
-		fmt.Println("set_audit_labels_invoked")
-
+	func(bctx rego.BuiltinContext, dataToStore *ast.Term) (*ast.Term, error) {
 		auditCache, err := GetAuditCache(bctx.Context)
 		if err != nil {
 			// TODO we should at least log a warning, consider failing the execution?
 			return nil, nil
 		}
+
 		data := make(map[string]interface{})
-		if err := ast.As(op2.Value, &data); err != nil {
+		if err := ast.As(dataToStore.Value, &data); err != nil {
 			return nil, err
 		}
-		fmt.Printf("ASDUTI CACHE %+v %+v\n", auditCache, data)
-		auditCache.Store("", data)
-		fmt.Printf("STORED\n")
 
+		auditCache.Store(data)
 		return ast.BooleanTerm(true), nil
 	},
 )

mocks/rego-policies/example.rego

@@ -93,7 +93,7 @@ assert_user {
 allow_policy_with_audit_data {
     print("HERE WE ARE")
     z := {"a": 42}
-    set_audit_labels("z", z)
+    set_audit_labels(z)
     true
 }
 projection_policy_with_audit_data[res] {