2 files changed, 67 insertions, 2 deletions
diff --git a/internal/flamegraph/nativesvg.go b/internal/flamegraph/nativesvg.go
index de0364b..26fc1e8 100644
--- a/internal/flamegraph/nativesvg.go
+++ b/internal/flamegraph/nativesvg.go
@@ -23,12 +23,17 @@ func NewNativeSVG(fields []string, countField string) NativeSVG {
 	}
 }
 
-func (n NativeSVG) WriteSVGFromFile(iorDataFile string) (string, error) {
-	outFile := fmt.Sprintf("%s.%s-by-%s.svg",
+func (n NativeSVG) WriteSVGFromFile(iorDataFile string) (outFile string, err error) {
+	outFile = fmt.Sprintf("%s.%s-by-%s.svg",
 		strings.TrimSuffix(iorDataFile, ".ior.zst"),
 		strings.Join(n.fields, ":"),
 		n.countField,
 	)
+	defer func() {
+		if err != nil {
+			_ = os.Remove(outFile)
+		}
+	}()
 
 	iod, err := newIorDataFromFile(iorDataFile)
 	if err != nil {
diff --git a/internal/flamegraph/nativesvg_test.go b/internal/flamegraph/nativesvg_test.go
new file mode 100644
index 0000000..36e88bf
--- /dev/null
+++ b/internal/flamegraph/nativesvg_test.go
@@ -0,0 +1,60 @@
+package flamegraph
+
+import (
+	"os"
+	"path/filepath"
+	"syscall"
+	"testing"
+
+	"ior/internal/types"
+
+	"github.com/DataDog/zstd"
+)
+
+func writeTestIorZst(t *testing.T, dir string) string {
+	t.Helper()
+
+	iod := newIorData()
+	iod.add("/tmp/test", types.SYS_ENTER_OPENAT, "tester", 100, 200, flagsType(syscall.O_RDONLY), Counter{
+		Count:          1,
+		Duration:       10,
+		DurationToPrev: 2,
+		Bytes:          0,
+	})
+	serialized, err := iod.serialize()
+	if err != nil {
+		t.Fatalf("serialize: %v", err)
+	}
+
+	path := filepath.Join(dir, "sample.ior.zst")
+	fd, err := os.Create(path)
+	if err != nil {
+		t.Fatalf("create test ior file: %v", err)
+	}
+	defer fd.Close()
+
+	enc := zstd.NewWriter(fd)
+	if _, err := enc.Write(serialized); err != nil {
+		t.Fatalf("write zstd payload: %v", err)
+	}
+	if err := enc.Close(); err != nil {
+		t.Fatalf("close zstd writer: %v", err)
+	}
+
+	return path
+}
+
+func TestWriteSVGFromFileCleansUpPartialOutputOnError(t *testing.T) {
+	dir := t.TempDir()
+	iorFile := writeTestIorZst(t, dir)
+
+	n := NewNativeSVG([]string{"invalidField"}, "count")
+	outFile, err := n.WriteSVGFromFile(iorFile)
+	if err == nil {
+		t.Fatal("expected error for invalid field, got nil")
+	}
+
+	if _, statErr := os.Stat(outFile); !os.IsNotExist(statErr) {
+		t.Fatalf("expected partial output to be removed, stat err=%v", statErr)
+	}
+}