github.com/waldiirawan/apm-agent-go/v2@v2.2.2/apmtest/recorder.go (about)

     1  // Licensed to Elasticsearch B.V. under one or more contributor
     2  // license agreements. See the NOTICE file distributed with
     3  // this work for additional information regarding copyright
     4  // ownership. Elasticsearch B.V. licenses this file to you under
     5  // the Apache License, Version 2.0 (the "License"); you may
     6  // not use this file except in compliance with the License.
     7  // You may obtain a copy of the License at
     8  //
     9  //     http://www.apache.org/licenses/LICENSE-2.0
    10  //
    11  // Unless required by applicable law or agreed to in writing,
    12  // software distributed under the License is distributed on an
    13  // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
    14  // KIND, either express or implied.  See the License for the
    15  // specific language governing permissions and limitations
    16  // under the License.
    17  
    18  package apmtest // import "github.com/waldiirawan/apm-agent-go/v2/apmtest"
    19  
    20  import (
    21  	"context"
    22  	"fmt"
    23  
    24  	"github.com/waldiirawan/apm-agent-go/v2"
    25  	"github.com/waldiirawan/apm-agent-go/v2/model"
    26  	"github.com/waldiirawan/apm-agent-go/v2/transport/transporttest"
    27  )
    28  
    29  // NewRecordingTracer returns a new RecordingTracer, containing a new
    30  // Tracer using the RecorderTransport stored inside.
    31  func NewRecordingTracer() *RecordingTracer {
    32  	var result RecordingTracer
    33  	tracer, err := apm.NewTracerOptions(apm.TracerOptions{
    34  		Transport: &result.RecorderTransport,
    35  	})
    36  	if err != nil {
    37  		panic(err)
    38  	}
    39  	result.Tracer = tracer
    40  	return &result
    41  }
    42  
    43  // RecordingTracer holds an apm.Tracer and transporttest.RecorderTransport.
    44  type RecordingTracer struct {
    45  	*apm.Tracer
    46  	transporttest.RecorderTransport
    47  }
    48  
    49  // WithTransaction calls rt.WithTransactionOptions with a zero apm.TransactionOptions.
    50  func (rt *RecordingTracer) WithTransaction(f func(ctx context.Context)) (model.Transaction, []model.Span, []model.Error) {
    51  	return rt.WithTransactionOptions(apm.TransactionOptions{}, f)
    52  }
    53  
    54  // WithTransactionOptions starts a transaction with the given options,
    55  // calls f with the transaction in the provided context, ends the transaction
    56  // and flushes the tracer, and then returns the resulting events.
    57  func (rt *RecordingTracer) WithTransactionOptions(opts apm.TransactionOptions, f func(ctx context.Context)) (model.Transaction, []model.Span, []model.Error) {
    58  	tx := rt.StartTransactionOptions("name", "type", opts)
    59  	ctx := apm.ContextWithTransaction(context.Background(), tx)
    60  	f(ctx)
    61  
    62  	tx.End()
    63  	rt.Flush(nil)
    64  	payloads := rt.Payloads()
    65  	if n := len(payloads.Transactions); n != 1 {
    66  		panic(fmt.Errorf("expected 1 transaction, got %d", n))
    67  	}
    68  	return payloads.Transactions[0], payloads.Spans, payloads.Errors
    69  }