schema/mixin: add AnnotateEdges helper (#659)

schema/mixin/mixin.go

@@ -8,6 +8,7 @@ import (
 	"time"
 
 	"github.com/facebookincubator/ent"
+	"github.com/facebookincubator/ent/schema/edge"
 	"github.com/facebookincubator/ent/schema/field"
 )
 
@@ -82,19 +83,38 @@ var _ ent.Mixin = (*Time)(nil)
 
 // AnnotateFields adds field annotations to underlying mixin fields.
 func AnnotateFields(m ent.Mixin, annotations ...field.Annotation) ent.Mixin {
-	return annotator{Mixin: m, annotations: annotations}
+	return fieldAnnotator{Mixin: m, annotations: annotations}
 }
 
-type annotator struct {
+// AnnotateEdges adds edge annotations to underlying mixin edges.
+func AnnotateEdges(m ent.Mixin, annotations ...edge.Annotation) ent.Mixin {
+	return edgeAnnotator{Mixin: m, annotations: annotations}
+}
+
+type fieldAnnotator struct {
 	ent.Mixin
 	annotations []field.Annotation
 }
 
-func (a annotator) Fields() []ent.Field {
+func (a fieldAnnotator) Fields() []ent.Field {
 	fields := a.Mixin.Fields()
-	for _, f := range fields {
-		desc := f.Descriptor()
+	for i := range fields {
+		desc := fields[i].Descriptor()
 		desc.Annotations = append(desc.Annotations, a.annotations...)
 	}
 	return fields
 }
+
+type edgeAnnotator struct {
+	ent.Mixin
+	annotations []edge.Annotation
+}
+
+func (a edgeAnnotator) Edges() []ent.Edge {
+	edges := a.Mixin.Edges()
+	for i := range edges {
+		desc := edges[i].Descriptor()
+		desc.Annotations = append(desc.Annotations, a.annotations...)
+	}
+	return edges
+}

schema/mixin/mixin_test.go

@@ -7,8 +7,11 @@ package mixin_test
 import (
 	"testing"
 
+	"github.com/facebookincubator/ent"
+	"github.com/facebookincubator/ent/schema/edge"
 	"github.com/facebookincubator/ent/schema/field"
 	"github.com/facebookincubator/ent/schema/mixin"
+
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -47,7 +50,7 @@ type annotation string
 
 func (annotation) Name() string { return "" }
 
-func TestAnnotate(t *testing.T) {
+func TestAnnotateFields(t *testing.T) {
 	annotations := []field.Annotation{
 		annotation("foo"),
 		annotation("bar"),
@@ -56,6 +59,7 @@ func TestAnnotate(t *testing.T) {
 	fields := mixin.AnnotateFields(
 		mixin.Time{}, annotations...,
 	).Fields()
+	require.Len(t, fields, 2)
 	for _, f := range fields {
 		desc := f.Descriptor()
 		require.Len(t, desc.Annotations, len(annotations))
@@ -64,3 +68,34 @@ func TestAnnotate(t *testing.T) {
 		}
 	}
 }
+
+type TestSchema struct {
+	ent.Schema
+}
+
+func (TestSchema) Edges() []ent.Edge {
+	return []ent.Edge{
+		edge.To("one", TestSchema.Type),
+		edge.From("two", TestSchema.Type).
+			Ref("one"),
+	}
+}
+
+func TestAnnotateEdges(t *testing.T) {
+	annotations := []edge.Annotation{
+		annotation("foo"),
+		annotation("bar"),
+		annotation("baz"),
+	}
+	edges := mixin.AnnotateEdges(
+		TestSchema{}, annotations...,
+	).Edges()
+	require.Len(t, edges, 2)
+	for _, e := range edges {
+		desc := e.Descriptor()
+		require.Len(t, desc.Annotations, len(annotations))
+		for i := range desc.Annotations {
+			assert.Equal(t, annotations[i], desc.Annotations[i])
+		}
+	}
+}