|
1 | 1 | package backend |
2 | 2 |
|
3 | 3 | import ( |
4 | | - "testing" |
5 | | - |
6 | 4 | "github.com/mudler/LocalAI/pkg/grpc/proto" |
| 5 | + |
| 6 | + . "github.com/onsi/ginkgo/v2" |
| 7 | + . "github.com/onsi/gomega" |
7 | 8 | ) |
8 | 9 |
|
9 | | -func TestDiarizationResultFromProto_NormalisesSpeakerLabels(t *testing.T) { |
10 | | - in := &proto.DiarizeResponse{ |
11 | | - Duration: 10.5, |
12 | | - Language: "en", |
13 | | - Segments: []*proto.DiarizeSegment{ |
14 | | - {Start: 0.0, End: 1.0, Speaker: "5", Text: "hi"}, |
15 | | - {Start: 1.0, End: 2.0, Speaker: "2"}, |
16 | | - {Start: 2.0, End: 3.5, Speaker: "5"}, |
17 | | - {Start: 3.5, End: 4.0, Speaker: ""}, // empty → coerced to "0" |
18 | | - }, |
19 | | - } |
| 10 | +var _ = Describe("diarizationResultFromProto", func() { |
| 11 | + It("normalises raw backend speaker labels to SPEAKER_NN in first-seen order", func() { |
| 12 | + in := &proto.DiarizeResponse{ |
| 13 | + Duration: 10.5, |
| 14 | + Language: "en", |
| 15 | + Segments: []*proto.DiarizeSegment{ |
| 16 | + {Start: 0.0, End: 1.0, Speaker: "5", Text: "hi"}, |
| 17 | + {Start: 1.0, End: 2.0, Speaker: "2"}, |
| 18 | + {Start: 2.0, End: 3.5, Speaker: "5"}, |
| 19 | + {Start: 3.5, End: 4.0, Speaker: ""}, // empty → coerced to "0" |
| 20 | + }, |
| 21 | + } |
20 | 22 |
|
21 | | - got := diarizationResultFromProto(in) |
| 23 | + got := diarizationResultFromProto(in) |
22 | 24 |
|
23 | | - if got.Task != "diarize" { |
24 | | - t.Errorf("task = %q, want diarize", got.Task) |
25 | | - } |
26 | | - if got.NumSpeakers != 3 { |
27 | | - t.Errorf("num_speakers = %d, want 3 (5, 2, 0)", got.NumSpeakers) |
28 | | - } |
29 | | - if got.Duration != 10.5 || got.Language != "en" { |
30 | | - t.Errorf("metadata not preserved: duration=%v language=%q", got.Duration, got.Language) |
31 | | - } |
32 | | - if len(got.Segments) != 4 { |
33 | | - t.Fatalf("segments=%d, want 4", len(got.Segments)) |
34 | | - } |
35 | | - // First-seen-order normalization: "5"→SPEAKER_00, "2"→SPEAKER_01, ""→SPEAKER_02 |
36 | | - want := []struct { |
37 | | - speaker string |
38 | | - label string |
39 | | - }{ |
40 | | - {"SPEAKER_00", "5"}, |
41 | | - {"SPEAKER_01", "2"}, |
42 | | - {"SPEAKER_00", "5"}, |
43 | | - {"SPEAKER_02", "0"}, |
44 | | - } |
45 | | - for i, w := range want { |
46 | | - if got.Segments[i].Speaker != w.speaker { |
47 | | - t.Errorf("seg[%d].speaker = %q, want %q", i, got.Segments[i].Speaker, w.speaker) |
| 25 | + Expect(got.Task).To(Equal("diarize")) |
| 26 | + Expect(got.NumSpeakers).To(Equal(3), "expected 3 distinct speakers (5, 2, 0)") |
| 27 | + Expect(got.Duration).To(BeEquivalentTo(10.5)) |
| 28 | + Expect(got.Language).To(Equal("en")) |
| 29 | + Expect(got.Segments).To(HaveLen(4)) |
| 30 | + |
| 31 | + // First-seen-order normalisation: "5"→SPEAKER_00, "2"→SPEAKER_01, ""→SPEAKER_02 |
| 32 | + want := []struct { |
| 33 | + speaker string |
| 34 | + label string |
| 35 | + }{ |
| 36 | + {"SPEAKER_00", "5"}, |
| 37 | + {"SPEAKER_01", "2"}, |
| 38 | + {"SPEAKER_00", "5"}, |
| 39 | + {"SPEAKER_02", "0"}, |
48 | 40 | } |
49 | | - if got.Segments[i].Label != w.label { |
50 | | - t.Errorf("seg[%d].label = %q, want %q", i, got.Segments[i].Label, w.label) |
| 41 | + for i, w := range want { |
| 42 | + Expect(got.Segments[i].Speaker).To(Equal(w.speaker), "seg[%d].speaker", i) |
| 43 | + Expect(got.Segments[i].Label).To(Equal(w.label), "seg[%d].label", i) |
51 | 44 | } |
52 | | - } |
53 | 45 |
|
54 | | - // Per-speaker totals reflect cumulative speech duration and segment count. |
55 | | - if len(got.Speakers) != 3 { |
56 | | - t.Fatalf("speakers=%d, want 3", len(got.Speakers)) |
57 | | - } |
58 | | - byID := map[string]float64{} |
59 | | - countByID := map[string]int{} |
60 | | - for _, sp := range got.Speakers { |
61 | | - byID[sp.Id] = sp.TotalSpeechDuration |
62 | | - countByID[sp.Id] = sp.SegmentCount |
63 | | - } |
64 | | - if byID["SPEAKER_00"] != 2.5 { // 1.0 + 1.5 |
65 | | - t.Errorf("SPEAKER_00 total = %v, want 2.5", byID["SPEAKER_00"]) |
66 | | - } |
67 | | - if byID["SPEAKER_01"] != 1.0 { |
68 | | - t.Errorf("SPEAKER_01 total = %v, want 1.0", byID["SPEAKER_01"]) |
69 | | - } |
70 | | - if countByID["SPEAKER_00"] != 2 || countByID["SPEAKER_01"] != 1 || countByID["SPEAKER_02"] != 1 { |
71 | | - t.Errorf("segment counts wrong: %v", countByID) |
72 | | - } |
73 | | -} |
| 46 | + // Per-speaker totals reflect cumulative speech duration and segment count. |
| 47 | + Expect(got.Speakers).To(HaveLen(3)) |
| 48 | + byID := map[string]float64{} |
| 49 | + countByID := map[string]int{} |
| 50 | + for _, sp := range got.Speakers { |
| 51 | + byID[sp.Id] = sp.TotalSpeechDuration |
| 52 | + countByID[sp.Id] = sp.SegmentCount |
| 53 | + } |
| 54 | + Expect(byID["SPEAKER_00"]).To(BeNumerically("~", 2.5, 0.001), "1.0 + 1.5") |
| 55 | + Expect(byID["SPEAKER_01"]).To(BeNumerically("~", 1.0, 0.001)) |
| 56 | + Expect(countByID["SPEAKER_00"]).To(Equal(2)) |
| 57 | + Expect(countByID["SPEAKER_01"]).To(Equal(1)) |
| 58 | + Expect(countByID["SPEAKER_02"]).To(Equal(1)) |
| 59 | + }) |
74 | 60 |
|
75 | | -func TestDiarizationResultFromProto_NilSafe(t *testing.T) { |
76 | | - got := diarizationResultFromProto(nil) |
77 | | - if got == nil || got.Segments == nil { |
78 | | - t.Fatalf("nil input must return a non-nil result with a non-nil segments slice") |
79 | | - } |
80 | | - if len(got.Segments) != 0 { |
81 | | - t.Errorf("nil input segments = %d, want 0", len(got.Segments)) |
82 | | - } |
83 | | -} |
| 61 | + It("returns a non-nil result with a non-nil segments slice for nil input", func() { |
| 62 | + got := diarizationResultFromProto(nil) |
| 63 | + Expect(got).ToNot(BeNil()) |
| 64 | + Expect(got.Segments).ToNot(BeNil()) |
| 65 | + Expect(got.Segments).To(BeEmpty()) |
| 66 | + }) |
84 | 67 |
|
85 | | -func TestDiarizationResultFromProto_EmptySegmentsKeepsBackendSpeakerCount(t *testing.T) { |
86 | | - // When the backend reports a non-zero NumSpeakers but no segments |
87 | | - // (early stop, silence-only audio after VAD trim, etc.), the result |
88 | | - // should still surface the backend's count rather than zero. |
89 | | - in := &proto.DiarizeResponse{NumSpeakers: 2, Duration: 5} |
90 | | - got := diarizationResultFromProto(in) |
91 | | - if got.NumSpeakers != 2 { |
92 | | - t.Errorf("num_speakers = %d, want 2", got.NumSpeakers) |
93 | | - } |
94 | | - if len(got.Segments) != 0 { |
95 | | - t.Errorf("expected no segments, got %d", len(got.Segments)) |
96 | | - } |
97 | | -} |
| 68 | + It("keeps the backend speaker count when no segments are returned", func() { |
| 69 | + // Backend reports a non-zero NumSpeakers but no segments (early stop, |
| 70 | + // silence-only audio after VAD trim). Surface the backend's count. |
| 71 | + in := &proto.DiarizeResponse{NumSpeakers: 2, Duration: 5} |
| 72 | + got := diarizationResultFromProto(in) |
| 73 | + Expect(got.NumSpeakers).To(Equal(2)) |
| 74 | + Expect(got.Segments).To(BeEmpty()) |
| 75 | + }) |
| 76 | +}) |
0 commit comments