From 31105f408dc21d07964e43b9b036f8dd501b7df6 Mon Sep 17 00:00:00 2001
From: Alex Snast <alexsn@fb.com>
Date: Mon, 20 Jul 2020 11:48:42 +0300
Subject: [PATCH] schema/mixin: adding mixin.Annotate helper which adds
 anootations to mixin fields (#625)

Signed-off-by: Alex Snast <alexsn@fb.com>
---
 schema/mixin/mixin.go      | 19 +++++++++++++++++++
 schema/mixin/mixin_test.go | 32 ++++++++++++++++++++++++++++----
 2 files changed, 47 insertions(+), 4 deletions(-)

diff --git a/schema/mixin/mixin.go b/schema/mixin/mixin.go
index e9b45fa83..860fc19eb 100644
--- a/schema/mixin/mixin.go
+++ b/schema/mixin/mixin.go
@@ -79,3 +79,22 @@ func (Time) Fields() []ent.Field {
 
 // time mixin must implement `Mixin` interface.
 var _ ent.Mixin = (*Time)(nil)
+
+// Annotate adds field annotations to underlying mixin fields.
+func Annotate(m ent.Mixin, annotations ...field.Annotation) ent.Mixin {
+	return annotator{Mixin: m, annotations: annotations}
+}
+
+type annotator struct {
+	ent.Mixin
+	annotations []field.Annotation
+}
+
+func (a annotator) Fields() []ent.Field {
+	fields := a.Mixin.Fields()
+	for _, f := range fields {
+		desc := f.Descriptor()
+		desc.Annotations = append(desc.Annotations, a.annotations...)
+	}
+	return fields
+}
diff --git a/schema/mixin/mixin_test.go b/schema/mixin/mixin_test.go
index 4caea504c..dfc9e5515 100644
--- a/schema/mixin/mixin_test.go
+++ b/schema/mixin/mixin_test.go
@@ -2,11 +2,13 @@
 // This source code is licensed under the Apache 2.0 license found
 // in the LICENSE file in the root directory of this source tree.
 
-package mixin
+package mixin_test
 
 import (
 	"testing"
 
+	"github.com/facebookincubator/ent/schema/field"
+	"github.com/facebookincubator/ent/schema/mixin"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -14,7 +16,7 @@ import (
 func TestTimeMixin(t *testing.T) {
 	t.Run("Create", func(t *testing.T) {
 		t.Parallel()
-		fields := CreateTime{}.Fields()
+		fields := mixin.CreateTime{}.Fields()
 		require.Len(t, fields, 1)
 		desc := fields[0].Descriptor()
 		assert.Equal(t, "create_time", desc.Name)
@@ -24,7 +26,7 @@ func TestTimeMixin(t *testing.T) {
 	})
 	t.Run("Update", func(t *testing.T) {
 		t.Parallel()
-		fields := UpdateTime{}.Fields()
+		fields := mixin.UpdateTime{}.Fields()
 		require.Len(t, fields, 1)
 		desc := fields[0].Descriptor()
 		assert.Equal(t, "update_time", desc.Name)
@@ -34,9 +36,31 @@ func TestTimeMixin(t *testing.T) {
 	})
 	t.Run("Compose", func(t *testing.T) {
 		t.Parallel()
-		fields := Time{}.Fields()
+		fields := mixin.Time{}.Fields()
 		require.Len(t, fields, 2)
 		assert.Equal(t, "create_time", fields[0].Descriptor().Name)
 		assert.Equal(t, "update_time", fields[1].Descriptor().Name)
 	})
 }
+
+type annotation string
+
+func (annotation) Name() string { return "" }
+
+func TestAnnotate(t *testing.T) {
+	annotations := []field.Annotation{
+		annotation("foo"),
+		annotation("bar"),
+		annotation("baz"),
+	}
+	fields := mixin.Annotate(
+		mixin.Time{}, annotations...,
+	).Fields()
+	for _, f := range fields {
+		desc := f.Descriptor()
+		require.Len(t, desc.Annotations, len(annotations))
+		for i := range desc.Annotations {
+			assert.Equal(t, annotations[i], desc.Annotations[i])
+		}
+	}
+}