test(flamegraph): cover height edge branches (0p)

3 files changed, 133 insertions, 0 deletions

diff --git a/internal/tui/flamegraph/controls_test.go b/internal/tui/flamegraph/controls_test.go
new file mode 100644
index 0000000..4a703fd
--- /dev/null
+++ b/internal/tui/flamegraph/controls_test.go
@@ -0,0 +1,86 @@
+package flamegraph
+
+import (
+	"errors"
+	"testing"
+
+	coreflamegraph "ior/internal/flamegraph"
+)
+
+type setHeightErrorTrie struct {
+	*coreflamegraph.LiveTrie
+	err error
+}
+
+func (t *setHeightErrorTrie) SetHeightField(string) error {
+	return t.err
+}
+
+type forcedHeightFieldTrie struct {
+	*coreflamegraph.LiveTrie
+	field string
+}
+
+func (t *forcedHeightFieldTrie) HeightField() string {
+	return t.field
+}
+
+func TestToggleHeightFieldErrorKeepsExistingState(t *testing.T) {
+	base := coreflamegraph.NewLiveTrie([]string{"comm", "path", "tracepoint"}, "count", "")
+	liveTrie := &setHeightErrorTrie{
+		LiveTrie: base,
+		err:      errors.New("set-height failed"),
+	}
+	m := NewModel(liveTrie)
+	m.snapshot = &snapshotNode{Name: "root", Total: 10}
+	m.frames = []tuiFrame{{Name: "root", Path: "root", Total: 10}}
+
+	m.toggleHeightField()
+
+	if got, want := m.heightField, ""; got != want {
+		t.Fatalf("heightField = %q, want %q on SetHeightField error", got, want)
+	}
+	if got, want := m.statusMessage, "Height toggle error: set-height failed"; got != want {
+		t.Fatalf("statusMessage = %q, want %q", got, want)
+	}
+	if m.snapshot == nil || len(m.frames) == 0 {
+		t.Fatalf("expected snapshot/layout state to remain intact on SetHeightField error")
+	}
+}
+
+func TestHeightFieldLabelReturnsRawValueForUnknownField(t *testing.T) {
+	m := NewModel(nil)
+	m.heightField = "custom-height"
+
+	if got, want := m.heightFieldLabel(), "custom-height"; got != want {
+		t.Fatalf("heightFieldLabel() = %q, want %q", got, want)
+	}
+}
+
+func TestSyncHeightFieldToTrieNilLiveTrieClearsField(t *testing.T) {
+	m := NewModel(nil)
+	m.heightField = "bytes"
+
+	m.syncHeightFieldToTrie()
+
+	if got, want := m.heightField, ""; got != want {
+		t.Fatalf("heightField = %q, want %q for nil liveTrie", got, want)
+	}
+}
+
+func TestSyncHeightFieldToTrieUnknownFieldFallsBackToOff(t *testing.T) {
+	base := coreflamegraph.NewLiveTrie([]string{"comm", "path", "tracepoint"}, "count", "")
+	liveTrie := &forcedHeightFieldTrie{
+		LiveTrie: base,
+		field:    "mystery",
+	}
+	m := NewModel(nil)
+	m.liveTrie = liveTrie
+	m.heightField = "duration"
+
+	m.syncHeightFieldToTrie()
+
+	if got, want := m.heightField, ""; got != want {
+		t.Fatalf("heightField = %q, want %q for unknown trie height field", got, want)
+	}
+}
diff --git a/internal/tui/flamegraph/model_test.go b/internal/tui/flamegraph/model_test.go
index 3efe589..cbbacef 100644
--- a/internal/tui/flamegraph/model_test.go
+++ b/internal/tui/flamegraph/model_test.go
@@ -355,6 +355,47 @@ func TestZoomLineageParentsAreNeverNarrowerThanChildren(t *testing.T) {
 	}
 }
 
+func TestApplyZoomLineagePreservesHeightTotals(t *testing.T) {
+	snapshot := &snapshotNode{
+		Name: "root",
+		Children: []*snapshotNode{
+			{
+				Name: "A",
+				Children: []*snapshotNode{
+					{Name: "A1", Total: 7, HeightTotal: 70},
+					{Name: "A2", Total: 3, HeightTotal: 30},
+				},
+			},
+			{Name: "B", Total: 4, HeightTotal: 40},
+		},
+	}
+	zoomPath := "root" + pathSeparator + "A"
+	zoomRoot := findNodeByPath(snapshot, zoomPath)
+	if zoomRoot == nil {
+		t.Fatalf("expected zoom root for %q", zoomPath)
+	}
+	zoomFrames := buildTerminalLayoutWithPath(zoomRoot, 100, 20, zoomPath)
+	zoomRootFrame := mustFindFrame(t, zoomFrames, zoomPath)
+	if got, want := zoomRootFrame.HeightTotal, uint64(100); got != want {
+		t.Fatalf("zoom root HeightTotal = %d, want %d before lineage", got, want)
+	}
+
+	frames := applyZoomLineage(zoomFrames, snapshot, zoomPath, 100)
+	root := mustFindFrame(t, frames, "root")
+	zoom := mustFindFrame(t, frames, zoomPath)
+	leaf := mustFindFrame(t, frames, zoomPath+pathSeparator+"A1")
+
+	if got, want := root.HeightTotal, uint64(140); got != want {
+		t.Fatalf("lineage root HeightTotal = %d, want %d", got, want)
+	}
+	if got, want := zoom.HeightTotal, zoomRootFrame.HeightTotal; got != want {
+		t.Fatalf("lineage zoom HeightTotal = %d, want %d", got, want)
+	}
+	if got, want := leaf.HeightTotal, uint64(70); got != want {
+		t.Fatalf("descendant HeightTotal = %d, want %d", got, want)
+	}
+}
+
 func TestMouseClickOnLineageAncestorUndoesToThatZoomLevel(t *testing.T) {
 	m := newZoomModel()
 	m.selectedIdx = mustFrameIndex(t, m.frames, "root"+pathSeparator+"A")
diff --git a/internal/tui/flamegraph/renderer_test.go b/internal/tui/flamegraph/renderer_test.go
index b078532..47f7d20 100644
--- a/internal/tui/flamegraph/renderer_test.go
+++ b/internal/tui/flamegraph/renderer_test.go
@@ -599,6 +599,12 @@ func TestBuildTerminalLayoutHeightTotalUsesSnapshotAggregation(t *testing.T) {
 	}
 }
 
+func TestSnapshotHeightTotalNilNodeReturnsZero(t *testing.T) {
+	if got, want := snapshotHeightTotal(nil), uint64(0); got != want {
+		t.Fatalf("snapshotHeightTotal(nil) = %d, want %d", got, want)
+	}
+}
+
 func mustFindFrame(t *testing.T, frames []tuiFrame, path string) tuiFrame {
 	t.Helper()
 	for _, frame := range frames {