From 6f2ce596acf9b9e32f740d887f7d0c8e9a469601 Mon Sep 17 00:00:00 2001
From: spinillos <selenepinillos@gmail.com>
Date: Thu, 20 Apr 2023 11:42:11 +0200
Subject: [PATCH] Modify txtar tests to add configurable import mappers

---
 generate_test.go               | 14 +++++++++----
 internal/cuetxtar/cuetxtar.go  | 12 +++++++++++
 testdata/imports/imports.txtar | 37 ++++++++++++++++++++++++++++++++++
 ts/ast/ast.go                  |  6 +++++-
 4 files changed, 64 insertions(+), 5 deletions(-)
 create mode 100644 testdata/imports/imports.txtar

diff --git a/generate_test.go b/generate_test.go
index de9f1d3..abb1887 100644
--- a/generate_test.go
+++ b/generate_test.go
@@ -46,6 +46,14 @@ func TestGenerateWithImports(t *testing.T) {
 		ToDo: map[string]string{
 			"imports/oneref_verbose": "Figure out how to disambiguate struct literals from the struct-with-braces-and-one-element case",
 		},
+		ImportMappers: map[string]func(string) (string, error){
+			"imports/imports": func(s string) (string, error) {
+				if s == "example.com/dep" {
+					return "new", nil
+				}
+				return s, nil
+			},
+		},
 	}
 
 	ctx := cuecontext.New()
@@ -57,10 +65,8 @@ func TestGenerateWithImports(t *testing.T) {
 		}
 
 		b, err := cuetsy.Generate(v, cuetsy.Config{
-			Export: true,
-			ImportMapper: func(string) (string, error) {
-				return "", nil
-			},
+			Export:       true,
+			ImportMapper: t.ImportMapperFn,
 		})
 		if err != nil {
 			errors.Print(t, err, nil)
diff --git a/internal/cuetxtar/cuetxtar.go b/internal/cuetxtar/cuetxtar.go
index 0325bc9..1a2abeb 100644
--- a/internal/cuetxtar/cuetxtar.go
+++ b/internal/cuetxtar/cuetxtar.go
@@ -59,6 +59,8 @@ type TxTarTest struct {
 
 	// ToDo is a map of tests that should be skipped now, but should be fixed.
 	ToDo map[string]string
+
+	ImportMappers map[string]func(string) (string, error)
 }
 
 // A Test represents a single test based on a .txtar file.
@@ -83,6 +85,8 @@ type Test struct {
 	Dir string
 
 	hasGold bool
+
+	ImportMapperFn func(s string) (string, error)
 }
 
 func (t *Test) Write(b []byte) (n int, err error) {
@@ -288,6 +292,14 @@ func (x *TxTarTest) Run(t *testing.T, f func(tc *Test)) {
 				t.Skip(msg)
 			}
 
+			if im, ok := x.ImportMappers[testName]; ok {
+				tc.ImportMapperFn = im
+			} else {
+				tc.ImportMapperFn = func(s string) (string, error) {
+					return "", nil
+				}
+			}
+
 			f(tc)
 
 			update := false
diff --git a/testdata/imports/imports.txtar b/testdata/imports/imports.txtar
new file mode 100644
index 0000000..db0efa6
--- /dev/null
+++ b/testdata/imports/imports.txtar
@@ -0,0 +1,37 @@
+-- cue.mod/module.cue --
+module: "example.com"
+
+-- in.cue --
+package test
+
+import "example.com/dep"
+import "example.com/other"
+
+#Foo: {
+	x: dep.#Test
+	y: other.#Test
+} @cuetsy(kind="interface")
+
+-- dep/file.cue --
+package dep
+
+#Test: {
+    a: string
+} @cuetsy(kind="interface")
+
+-- other/file.cue --
+package other
+
+#Test: {
+    a: string
+} @cuetsy(kind="interface")
+
+-- out/gen --
+
+import * as dep from 'new';
+import * as other from 'example.com/other';
+
+export interface Foo {
+  x: dep.Test;
+  y: other.Test;
+}
diff --git a/ts/ast/ast.go b/ts/ast/ast.go
index 619971f..d0f61d7 100644
--- a/ts/ast/ast.go
+++ b/ts/ast/ast.go
@@ -49,7 +49,11 @@ func (f File) String() string {
 
 	for _, i := range f.Imports {
 		b.WriteString(formatInner(EOLNone, 0, i))
-		b.WriteString("\n\n")
+		b.WriteString("\n")
+	}
+
+	if len(f.Imports) > 0 {
+		b.WriteString("\n")
 	}
 
 	for i, n := range f.Nodes {