Add unit tests for anonymizer/app/writer

cmd/anonymizer/app/writer/writer.go

@@ -17,6 +17,7 @@ package writer
 import (
 	"bytes"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"os"
 	"sync"
@@ -28,6 +29,8 @@ import (
 	"github.com/jaegertracing/jaeger/model"
 )
 
+var ErrMaxSpansCountReached = errors.New("max spans count reached")
+
 // Config contains parameters to NewWriter.
 type Config struct {
 	MaxSpansCount  int                `yaml:"max_spans_count" name:"max_spans_count"`
@@ -130,7 +133,7 @@ func (w *Writer) WriteSpan(msg *model.Span) error {
 	if w.config.MaxSpansCount > 0 && w.spanCount >= w.config.MaxSpansCount {
 		w.logger.Info("Saved enough spans, exiting...")
 		w.Close()
-		os.Exit(0)
+		return ErrMaxSpansCountReached
 	}
 
 	return nil

cmd/anonymizer/app/writer/writer_test.go

@@ -0,0 +1,117 @@
+// Copyright (c) 2024 The Jaeger Authors.
+// SPDX-License-Identifier: Apache-2.0
+
+package writer
+
+import (
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+	"go.uber.org/zap"
+
+	"github.com/jaegertracing/jaeger/model"
+)
+
+var tags = []model.KeyValue{
+	model.Bool("error", true),
+	model.String("http.method", "POST"),
+	model.Bool("foobar", true),
+}
+
+var traceID = model.NewTraceID(1, 2)
+
+var span = &model.Span{
+	TraceID: traceID,
+	SpanID:  model.NewSpanID(1),
+	Process: &model.Process{
+		ServiceName: "serviceName",
+		Tags:        tags,
+	},
+	OperationName: "operationName",
+	Tags:          tags,
+	Logs: []model.Log{
+		{
+			Timestamp: time.Now(),
+			Fields: []model.KeyValue{
+				model.String("logKey", "logValue"),
+			},
+		},
+	},
+	Duration:  time.Second * 5,
+	StartTime: time.Unix(300, 0),
+}
+
+func TestNew(t *testing.T) {
+	nopLogger := zap.NewNop()
+	tempDir := t.TempDir()
+
+	t.Run("no error", func(t *testing.T) {
+		config := Config{
+			MaxSpansCount:  10,
+			CapturedFile:   tempDir + "/captured.json",
+			AnonymizedFile: tempDir + "/anonymized.json",
+			MappingFile:    tempDir + "/mapping.json",
+		}
+		_, err := New(config, nopLogger)
+		require.NoError(t, err)
+	})
+
+	t.Run("CapturedFile does not exist", func(t *testing.T) {
+		config := Config{
+			CapturedFile:   tempDir + "/nonexistent_directory/captured.json",
+			AnonymizedFile: tempDir + "/anonymized.json",
+			MappingFile:    tempDir + "/mapping.json",
+		}
+		_, err := New(config, nopLogger)
+		require.ErrorContains(t, err, "cannot create output file")
+	})
+
+	t.Run("AnonymizedFile does not exist", func(t *testing.T) {
+		config := Config{
+			CapturedFile:   tempDir + "/captured.json",
+			AnonymizedFile: tempDir + "/nonexistent_directory/anonymized.json",
+			MappingFile:    tempDir + "/mapping.json",
+		}
+		_, err := New(config, nopLogger)
+		require.ErrorContains(t, err, "cannot create output file")
+	})
+}
+
+func TestWriter_WriteSpan(t *testing.T) {
+	nopLogger := zap.NewNop()
+	t.Run("write span", func(t *testing.T) {
+		tempDir := t.TempDir()
+		config := Config{
+			MaxSpansCount:  10,
+			CapturedFile:   tempDir + "/captured.json",
+			AnonymizedFile: tempDir + "/anonymized.json",
+			MappingFile:    tempDir + "/mapping.json",
+		}
+
+		writer, err := New(config, nopLogger)
+		require.NoError(t, err)
+
+		for i := 0; i < 9; i++ {
+			err = writer.WriteSpan(span)
+			require.NoError(t, err)
+		}
+		defer writer.Close()
+	})
+	t.Run("write span with MaxSpansCount", func(t *testing.T) {
+		tempDir := t.TempDir()
+		config := Config{
+			MaxSpansCount:  1,
+			CapturedFile:   tempDir + "/captured.json",
+			AnonymizedFile: tempDir + "/anonymized.json",
+			MappingFile:    tempDir + "/mapping.json",
+		}
+
+		writer, err := New(config, zap.NewNop())
+		require.NoError(t, err)
+
+		err = writer.WriteSpan(span)
+		require.ErrorContains(t, err, "saved MaxSpansCount")
+		defer writer.Close()
+	})
+}

cmd/anonymizer/main.go

@@ -15,6 +15,7 @@
 package main
 
 import (
+	"errors"
 	"fmt"
 	"os"
 
@@ -53,7 +54,7 @@ func main() {
 				},
 			}
 
-			writer, err := writer.New(conf, logger)
+			writerObj, err := writer.New(conf, logger)
 			if err != nil {
 				logger.Fatal("error while creating writer object", zap.Error(err))
 			}
@@ -69,9 +70,15 @@ func main() {
 			}
 
 			for _, span := range spans {
-				writer.WriteSpan(&span)
+				err := writerObj.WriteSpan(&span)
+				if err != nil {
+					logger.Fatal("error while writing span", zap.Error(err))
+					if errors.Is(err, writer.ErrMaxSpansCountReached) {
+						os.Exit(0)
+					}
+				}
 			}
-			writer.Close()
+			writerObj.Close()
 
 			uiCfg := uiconv.Config{
 				CapturedFile: conf.AnonymizedFile,