From eea4ec78205e955b757fbb3166cd274618c00742 Mon Sep 17 00:00:00 2001
From: sam boyer <sdboyer@grafana.com>
Date: Fri, 21 Jul 2023 17:25:22 -0400
Subject: [PATCH] Tests for other bad inputs, add new translate impl

---
 gomig_test.go | 58 +++++++++++++++++++++++++++++++++++++++++++++++++--
 instance.go   | 38 +++++++++++++++++++++++++++++++++
 lineage.go    |  3 ++-
 3 files changed, 96 insertions(+), 3 deletions(-)

diff --git a/gomig_test.go b/gomig_test.go
index 109c692..463b6e5 100644
--- a/gomig_test.go
+++ b/gomig_test.go
@@ -3,6 +3,8 @@ package thema
 import (
 	"testing"
 
+	"github.com/stretchr/testify/assert"
+
 	"cuelang.org/go/cue/cuecontext"
 	"github.com/stretchr/testify/require"
 )
@@ -141,7 +143,7 @@ schemas: [{
 }]
 `
 
-	lenses := []ImperativeLens{
+	correctLenses := []ImperativeLens{
 		{
 			To:   SV(0, 0),
 			From: SV(0, 1),
@@ -313,8 +315,60 @@ schemas: [{
 	ctx := cuecontext.New()
 	rt := NewRuntime(ctx)
 	linval := rt.Context().CompileString(multivlin)
-	_, err := BindLineage(linval, rt, ImperativeLenses(lenses...))
+	_, err := BindLineage(linval, rt, ImperativeLenses(correctLenses...))
 	require.NoError(t, err)
+
+	_, err = BindLineage(linval, rt, ImperativeLenses(correctLenses[1:]...))
+	assert.Error(t, err, "expected error when missing a reverse Go migration")
+	_, err = BindLineage(linval, rt, ImperativeLenses(correctLenses[:1]...))
+	assert.Error(t, err, "expected error when missing a forward Go migration")
+
+	_, err = BindLineage(linval, rt, ImperativeLenses(append(correctLenses, ImperativeLens{
+		To:     SV(2, 0),
+		From:   SV(2, 1),
+		Mapper: func(inst *Instance, to Schema) (*Instance, error) { return nil, nil },
+	})...))
+	assert.Error(t, err, "expected error when adding Go migration pointing to nonexistent version")
+
+	_, err = BindLineage(linval, rt, ImperativeLenses(append(correctLenses, ImperativeLens{
+		To:     SV(2, 1),
+		From:   SV(2, 0),
+		Mapper: func(inst *Instance, to Schema) (*Instance, error) { return nil, nil },
+	})...))
+	assert.Error(t, err, "expected error when adding Go migration pointing to nonexistent version")
+
+	_, err = BindLineage(linval, rt, ImperativeLenses(append(correctLenses, ImperativeLens{
+		To:     SV(2, 0),
+		From:   SV(1, 1),
+		Mapper: func(inst *Instance, to Schema) (*Instance, error) { return nil, nil },
+	})...))
+	assert.Error(t, err, "expected error when adding duplicate Go migration")
+
+	_, err = BindLineage(linval, rt, ImperativeLenses(append(correctLenses, ImperativeLens{
+		Mapper: func(inst *Instance, to Schema) (*Instance, error) { return nil, nil },
+	})...))
+	assert.Error(t, err, "expected error when providing a Go migration with same to and from")
+
+	_, err = BindLineage(linval, rt, ImperativeLenses(append(correctLenses, ImperativeLens{
+		To:     SV(2, 0),
+		From:   SV(1, 0),
+		Mapper: func(inst *Instance, to Schema) (*Instance, error) { return nil, nil },
+	})...))
+	assert.Error(t, err, "expected error when providing Go migration with wrong successor")
+
+	_, err = BindLineage(linval, rt, ImperativeLenses(append(correctLenses, ImperativeLens{
+		To:     SV(1, 0),
+		From:   SV(2, 0),
+		Mapper: func(inst *Instance, to Schema) (*Instance, error) { return nil, nil },
+	})...))
+	assert.Error(t, err, "expected error when providing Go migration with wrong predecessor")
+
+	_, err = BindLineage(linval, rt, ImperativeLenses(append(correctLenses, ImperativeLens{
+		To:     SV(1, 1),
+		From:   SV(1, 0),
+		Mapper: func(inst *Instance, to Schema) (*Instance, error) { return nil, nil },
+	})...))
+	assert.Error(t, err, "expected error when providing Go migration for minor version upgrade")
 }
 
 func tomap(inst *Instance) map[string]any {
diff --git a/instance.go b/instance.go
index 4a5c3bf..9381f16 100644
--- a/instance.go
+++ b/instance.go
@@ -198,6 +198,10 @@ func (inst *TypedInstance[T]) ValueP() T {
 func (i *Instance) Translate(to SyntacticVersion) (*Instance, TranslationLacunas, error) {
 	i.check()
 
+	if len(i.Schema().Lineage().(*baseLineage).lensmap) > 0 {
+		return i.translateGo(to)
+	}
+
 	// TODO define this in terms of AsSuccessor and AsPredecessor, rather than those in terms of this.
 	newsch, err := i.Schema().Lineage().Schema(to)
 	if err != nil {
@@ -240,6 +244,40 @@ func (i *Instance) Translate(to SyntacticVersion) (*Instance, TranslationLacunas
 	return inst, lac, err
 }
 
+func (i *Instance) translateGo(to SyntacticVersion) (*Instance, TranslationLacunas, error) {
+	from := i.Schema().Version()
+	if to == from {
+		// TODO make sure this mirrors the pure CUE behavior
+		return i, nil, nil
+	}
+	lensmap := i.Schema().Lineage().(*baseLineage).lensmap
+
+	sch := i.Schema()
+	ti := new(Instance)
+	*ti = *i
+	for sch.Version() != to {
+		rti, err := lensmap[lid(sch.Version(), ti.Schema().Version())].Mapper(ti, sch)
+		if err != nil {
+			return nil, nil, err
+		}
+		// Ensure the returned instance exists and the caller returned an instance of the expected schema version
+		if rti == nil {
+			return nil, nil, fmt.Errorf("lens returned a nil instance")
+		}
+		if rti.Schema().Version() != sch.Version() {
+			return nil, nil, fmt.Errorf("lens returned an instance of the wrong schema version: expected %v, got %v", sch.Version(), ti.Schema().Version())
+		}
+		*ti = *rti
+		if to.Less(from) {
+			sch = sch.Predecessor()
+		} else {
+			sch = sch.Successor()
+		}
+	}
+
+	return ti, nil, nil
+}
+
 type multiTranslationLacunas []struct {
 	V   SyntacticVersion `json:"v"`
 	Lac []Lacuna         `json:"lacunas"`
diff --git a/lineage.go b/lineage.go
index 66c2019..6511269 100644
--- a/lineage.go
+++ b/lineage.go
@@ -40,7 +40,7 @@ type baseLineage struct {
 	// all the schemas
 	allsch []*schemaDef
 
-	implens []ImperativeLens
+	lensmap map[lensID]ImperativeLens
 }
 
 // BindLineage takes a raw [cue.Value], checks that it correctly follows Thema's
@@ -143,6 +143,7 @@ func BindLineage(v cue.Value, rt *Runtime, opts ...BindOption) (Lineage, error)
 		uni:       ml.uni,
 		allsch:    ml.schlist,
 		allv:      ml.allv,
+		lensmap:   ml.lensmap,
 	}
 
 	for _, sch := range lin.allsch {