fix: return errors for unknown counter fields (task 383)

3 files changed, 37 insertions, 14 deletions

diff --git a/internal/flamegraph/counter.go b/internal/flamegraph/counter.go
index 539d017..441db68 100644
--- a/internal/flamegraph/counter.go
+++ b/internal/flamegraph/counter.go
@@ -28,17 +28,17 @@ func (c Counter) add(other Counter) Counter {
 	return c
 }
 
-func (c Counter) ValueByName(name string) uint64 {
+func (c Counter) ValueByName(name string) (uint64, error) {
 	switch name {
 	case "count":
-		return c.Count
+		return c.Count, nil
 	case "duration":
-		return c.Duration
+		return c.Duration, nil
 	case "durationToPrev":
-		return c.DurationToPrev
+		return c.DurationToPrev, nil
 	case "bytes":
-		return c.Bytes
+		return c.Bytes, nil
 	default:
-		panic(fmt.Sprintln("No", name, "in count record"))
+		return 0, fmt.Errorf("unknown counter field %q", name)
 	}
 }
diff --git a/internal/flamegraph/iordata_test.go b/internal/flamegraph/iordata_test.go
index f4855b0..ee07a90 100644
--- a/internal/flamegraph/iordata_test.go
+++ b/internal/flamegraph/iordata_test.go
@@ -170,16 +170,36 @@ func TestStringByNameValidFields(t *testing.T) {
 	}
 }
 
-func TestCounterValueByNamePanic(t *testing.T) {
+func TestCounterValueByNameUnknownField(t *testing.T) {
 	c := Counter{Count: 1, Duration: 100, DurationToPrev: 10, Bytes: 64}
 
-	defer func() {
-		if r := recover(); r == nil {
-			t.Error("Expected panic for unknown counter name, got none")
-		}
-	}()
+	_, err := c.ValueByName("nonexistent")
+	if err == nil {
+		t.Error("Expected error for unknown counter name, got nil")
+	}
+}
 
-	c.ValueByName("nonexistent")
+func TestCounterValueByNameValidFields(t *testing.T) {
+	c := Counter{Count: 1, Duration: 100, DurationToPrev: 10, Bytes: 64}
+
+	tests := map[string]uint64{
+		"count":          c.Count,
+		"duration":       c.Duration,
+		"durationToPrev": c.DurationToPrev,
+		"bytes":          c.Bytes,
+	}
+
+	for field, want := range tests {
+		t.Run(field, func(t *testing.T) {
+			got, err := c.ValueByName(field)
+			if err != nil {
+				t.Fatalf("Expected no error for field %q, got %v", field, err)
+			}
+			if got != want {
+				t.Fatalf("Expected %d for field %q, got %d", want, field, got)
+			}
+		})
+	}
 }
 
 func TestMergeEmpty(t *testing.T) {
diff --git a/internal/flamegraph/livetrie.go b/internal/flamegraph/livetrie.go
index 04137bf..e76a806 100644
--- a/internal/flamegraph/livetrie.go
+++ b/internal/flamegraph/livetrie.go
@@ -102,7 +102,10 @@ func (lt *LiveTrie) Ingest(ep *event.Pair) {
 
 // AddRecord adds one already-decoded flamegraph record into the live trie.
 func (lt *LiveTrie) AddRecord(record IterRecord) {
-	value := record.Cnt.ValueByName(lt.countField)
+	value, err := record.Cnt.ValueByName(lt.countField)
+	if err != nil {
+		return
+	}
 
 	lt.mu.Lock()
 	frames := lt.buildFrames(record)