diff --git a/cmd/anonymizer/app/writer/.nocover b/cmd/anonymizer/app/writer/.nocover deleted file mode 100644 index 5d2db86d7c0..00000000000 --- a/cmd/anonymizer/app/writer/.nocover +++ /dev/null @@ -1 +0,0 @@ -nobn-critical test utility diff --git a/cmd/anonymizer/app/writer/writer.go b/cmd/anonymizer/app/writer/writer.go index dee49e1735d..5b22a371c4b 100644 --- a/cmd/anonymizer/app/writer/writer.go +++ b/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 diff --git a/cmd/anonymizer/app/writer/writer_test.go b/cmd/anonymizer/app/writer/writer_test.go new file mode 100644 index 00000000000..73391682e30 --- /dev/null +++ b/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) + defer writer.Close() + + for i := 0; i < 9; i++ { + err = writer.WriteSpan(span) + require.NoError(t, err) + } + }) + 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) + defer writer.Close() + + err = writer.WriteSpan(span) + require.ErrorIs(t, err, ErrMaxSpansCountReached) + }) +} diff --git a/cmd/anonymizer/main.go b/cmd/anonymizer/main.go index 259321247b0..0e2949146f1 100644 --- a/cmd/anonymizer/main.go +++ b/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) + w, 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) + if err := w.WriteSpan(&span); err != nil { + if errors.Is(err, writer.ErrMaxSpansCountReached) { + logger.Info("max spans count reached") + os.Exit(0) + } + logger.Error("error while writing span", zap.Error(err)) + } } - writer.Close() + w.Close() uiCfg := uiconv.Config{ CapturedFile: conf.AnonymizedFile,