MockRecorder Total*Tokens

Author: pawbana

internal/testutil/mock_recorder.go

@@ -76,17 +76,21 @@ func (m *MockRecorder) RecordedTokenUsages() []*recorder.TokenUsageRecord {
 
 // TotalInputTokens returns the sum of input tokens across all recorded token usages.
 func (m *MockRecorder) TotalInputTokens() int64 {
+	m.mu.Lock()
+	defer m.mu.Unlock()
 	var total int64
-	for _, el := range m.RecordedTokenUsages() {
+	for _, el := range m.tokenUsages {
 		total += el.Input
 	}
 	return total
 }
 
 // TotalOutputTokens returns the sum of output tokens across all recorded token usages.
 func (m *MockRecorder) TotalOutputTokens() int64 {
+	m.mu.Lock()
+	defer m.mu.Unlock()
 	var total int64
-	for _, el := range m.RecordedTokenUsages() {
+	for _, el := range m.tokenUsages {
 		total += el.Output
 	}
 	return total