github.com/v2fly/tools@v0.100.0/internal/event/export/eventtest/eventtest.go (about)

     1  // Copyright 2020 The Go Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  // Package eventtest supports logging events to a test.
     6  // You can use NewContext to create a context that knows how to deliver
     7  // telemetry events back to the test.
     8  // You must use this context or a derived one anywhere you want telemetry to be
     9  // correctly routed back to the test it was constructed with.
    10  // Any events delivered to a background context will be dropped.
    11  //
    12  // Importing this package will cause it to register a new global telemetry
    13  // exporter that understands the special contexts returned by NewContext.
    14  // This means you should not import this package if you are not going to call
    15  // NewContext.
    16  package eventtest
    17  
    18  import (
    19  	"bytes"
    20  	"context"
    21  	"sync"
    22  	"testing"
    23  
    24  	"github.com/v2fly/tools/internal/event"
    25  	"github.com/v2fly/tools/internal/event/core"
    26  	"github.com/v2fly/tools/internal/event/export"
    27  	"github.com/v2fly/tools/internal/event/label"
    28  )
    29  
    30  func init() {
    31  	e := &testExporter{buffer: &bytes.Buffer{}}
    32  	e.logger = export.LogWriter(e.buffer, false)
    33  
    34  	event.SetExporter(export.Spans(e.processEvent))
    35  }
    36  
    37  type testingKeyType int
    38  
    39  const testingKey = testingKeyType(0)
    40  
    41  // NewContext returns a context you should use for the active test.
    42  func NewContext(ctx context.Context, t testing.TB) context.Context {
    43  	return context.WithValue(ctx, testingKey, t)
    44  }
    45  
    46  type testExporter struct {
    47  	mu     sync.Mutex
    48  	buffer *bytes.Buffer
    49  	logger event.Exporter
    50  }
    51  
    52  func (w *testExporter) processEvent(ctx context.Context, ev core.Event, tm label.Map) context.Context {
    53  	w.mu.Lock()
    54  	defer w.mu.Unlock()
    55  	// build our log message in buffer
    56  	result := w.logger(ctx, ev, tm)
    57  	v := ctx.Value(testingKey)
    58  	// get the testing.TB
    59  	if w.buffer.Len() > 0 && v != nil {
    60  		v.(testing.TB).Log(w.buffer)
    61  	}
    62  	w.buffer.Truncate(0)
    63  	return result
    64  }