More testing for ast, json

Signed-off-by: Ashutosh Narkar <anarkar4387@gmail.com>

ast/marshal_test.go

@@ -444,6 +444,22 @@ func TestRule_MarshalJSON(t *testing.T) {
 			}(),
 			ExpectedJSON: `{"body":[{"index":0,"terms":{"type":"boolean","value":true}}],"head":{"name":"allow","value":{"type":"boolean","value":true},"ref":[{"type":"var","value":"allow"}]},"location":{"file":"example.rego","row":6,"col":2}}`,
 		},
+		"annotations included": {
+			Rule: func() *Rule {
+				r := rule.Copy()
+				r.Annotations = []*Annotations{{
+					Scope:         "rule",
+					Title:         "My rule",
+					Entrypoint:    true,
+					Organizations: []string{"org1"},
+					Description:   "My desc",
+					Custom: map[string]interface{}{
+						"foo": "bar",
+					}}}
+				return r
+			}(),
+			ExpectedJSON: `{"annotations":[{"custom":{"foo":"bar"},"description":"My desc","entrypoint":true,"organizations":["org1"],"scope":"rule","title":"My rule"}],"body":[{"index":0,"terms":{"type":"boolean","value":true}}],"head":{"name":"allow","value":{"type":"boolean","value":true},"ref":[{"type":"var","value":"allow"}]}}`,
+		},
 	}
 
 	for name, data := range testCases {

ast/policy.go

@@ -658,6 +658,11 @@ func (rule *Rule) Compare(other *Rule) int {
 	if cmp := rule.Body.Compare(other.Body); cmp != 0 {
 		return cmp
 	}
+
+	if cmp := annotationsCompare(rule.Annotations, other.Annotations); cmp != 0 {
+		return cmp
+	}
+
 	return rule.Else.Compare(other.Else)
 }
 
@@ -764,6 +769,10 @@ func (rule *Rule) MarshalJSON() ([]byte, error) {
 		}
 	}
 
+	if len(rule.Annotations) != 0 {
+		data["annotations"] = rule.Annotations
+	}
+
 	return json.Marshal(data)
 }
 

compile/compile.go

@@ -1019,22 +1019,22 @@ func (o *optimizer) Do(ctx context.Context) error {
 		}
 
 		if len(pq.Support) != 0 {
-
 			// attach annotations to the module
 			for _, ar := range flattenedAnnotations {
 				for _, module := range pq.Support {
 					if module.Package.Path.Equal(ar.Path) {
 						module.Annotations = append(module.Annotations, ar.Annotations)
 					}
-
-					var annotations []*ast.Annotations
-					for _, rule := range module.Rules {
-						annotations = append(annotations, rule.Annotations...)
-					}
-					module.Annotations = append(module.Annotations, annotations...)
 				}
 			}
 
+			for _, module := range pq.Support {
+				var annotations []*ast.Annotations
+				for _, rule := range module.Rules {
+					annotations = append(annotations, rule.Annotations...)
+				}
+				module.Annotations = append(module.Annotations, annotations...)
+			}
 		}
 
 		if module := o.getSupportForEntrypoint(pq.Queries, e, resultsym, flattenedAnnotations); module != nil {
@@ -1138,6 +1138,13 @@ func (o *optimizer) getSupportForEntrypoint(queries []ast.Body, e *ast.Term, res
 	name := ast.Var(path[len(path)-1].Value.(ast.String))
 	module := &ast.Module{Package: &ast.Package{Path: path[:len(path)-1]}}
 
+	// attach annotations to the module
+	for _, ar := range annotationRefs {
+		if module.Package.Path.Equal(ar.Path) {
+			module.Annotations = append(module.Annotations, ar.Annotations)
+		}
+	}
+
 	ruleAnnotations := findAnnotationsForTerm(e, annotationRefs)
 
 	for _, query := range queries {
@@ -1166,21 +1173,9 @@ func (o *optimizer) getSupportForEntrypoint(queries []ast.Body, e *ast.Term, res
 		}
 
 		module.Rules = append(module.Rules, rule)
+		module.Annotations = append(module.Annotations, ruleAnnotations...)
 	}
 
-	// attach annotations to the module
-	var annotations []*ast.Annotations
-
-	for _, ar := range annotationRefs {
-		if module.Package.Path.Equal(ar.Path) {
-			annotations = append(annotations, ar.Annotations)
-		}
-	}
-
-	annotations = append(annotations, ruleAnnotations...)
-
-	module.Annotations = annotations
-
 	return module
 }