From a898e5021164bef9bcabc91848935911b1438279 Mon Sep 17 00:00:00 2001
From: Alexander Wang <alex@terrastruct.com>
Date: Wed, 2 Aug 2023 11:38:53 -0700
Subject: [PATCH] fmt test

---
 d2cli/fmt.go              | 10 ++++++++++
 e2etests-cli/main_test.go | 20 +++++++++++++++++++-
 2 files changed, 29 insertions(+), 1 deletion(-)

diff --git a/d2cli/fmt.go b/d2cli/fmt.go
index 5cb7999cb..d7aa62ca2 100644
--- a/d2cli/fmt.go
+++ b/d2cli/fmt.go
@@ -3,6 +3,8 @@ package d2cli
 import (
 	"bytes"
 	"context"
+	"os"
+	"path/filepath"
 
 	"oss.terrastruct.com/util-go/xdefer"
 
@@ -21,6 +23,14 @@ func fmtCmd(ctx context.Context, ms *xmain.State) (err error) {
 	}
 
 	for _, inputPath := range ms.Opts.Args {
+		if inputPath != "-" {
+			inputPath = ms.AbsPath(inputPath)
+			d, err := os.Stat(inputPath)
+			if err == nil && d.IsDir() {
+				inputPath = filepath.Join(inputPath, "index.d2")
+			}
+		}
+
 		input, err := ms.ReadPath(inputPath)
 		if err != nil {
 			return err
diff --git a/e2etests-cli/main_test.go b/e2etests-cli/main_test.go
index d955e351f..bd0385bd5 100644
--- a/e2etests-cli/main_test.go
+++ b/e2etests-cli/main_test.go
@@ -521,6 +521,16 @@ i used to read
 				assert.Testdata(t, ".svg", svg)
 			},
 		},
+		{
+			name: "basic-fmt",
+			run: func(t *testing.T, ctx context.Context, dir string, env *xos.Env) {
+				writeFile(t, dir, "hello-world.d2", `x ---> y`)
+				err := runTestMainPersist(t, ctx, dir, env, "fmt", "hello-world.d2")
+				assert.Success(t, err)
+				got := readFile(t, dir, "hello-world.d2")
+				assert.Equal(t, "x -> y\n", string(got))
+			},
+		},
 	}
 
 	ctx := context.Background()
@@ -561,6 +571,15 @@ func testMain(dir string, env *xos.Env, args ...string) *xmain.TestState {
 }
 
 func runTestMain(tb testing.TB, ctx context.Context, dir string, env *xos.Env, args ...string) error {
+	err := runTestMainPersist(tb, ctx, dir, env, args...)
+	if err != nil {
+		return err
+	}
+	removeD2Files(tb, dir)
+	return nil
+}
+
+func runTestMainPersist(tb testing.TB, ctx context.Context, dir string, env *xos.Env, args ...string) error {
 	tms := testMain(dir, env, args...)
 	tms.Start(tb, ctx)
 	defer tms.Cleanup(tb)
@@ -568,7 +587,6 @@ func runTestMain(tb testing.TB, ctx context.Context, dir string, env *xos.Env, a
 	if err != nil {
 		return err
 	}
-	removeD2Files(tb, dir)
 	return nil
 }