github.com/cloudwego/hertz@v0.9.3/internal/stats/tracer_test.go (about)

     1  /*
     2   * Copyright 2022 CloudWeGo Authors
     3   *
     4   * Licensed under the Apache License, Version 2.0 (the "License");
     5   * you may not use this file except in compliance with the License.
     6   * You may obtain a copy of the License at
     7   *
     8   *     http://www.apache.org/licenses/LICENSE-2.0
     9   *
    10   * Unless required by applicable law or agreed to in writing, software
    11   * distributed under the License is distributed on an "AS IS" BASIS,
    12   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13   * See the License for the specific language governing permissions and
    14   * limitations under the License.
    15   */
    16  
    17  package stats
    18  
    19  import (
    20  	"context"
    21  	"errors"
    22  	"fmt"
    23  	"testing"
    24  
    25  	"github.com/cloudwego/hertz/pkg/app"
    26  	"github.com/cloudwego/hertz/pkg/common/test/assert"
    27  	"github.com/cloudwego/hertz/pkg/common/tracer/traceinfo"
    28  )
    29  
    30  type mockTracer struct {
    31  	order         int
    32  	stack         *[]int
    33  	panicAtStart  bool
    34  	panicAtFinish bool
    35  }
    36  
    37  func (mt *mockTracer) Start(ctx context.Context, c *app.RequestContext) context.Context {
    38  	if mt.panicAtStart {
    39  		panic(fmt.Sprintf("panicked at start: Tracer(%d)", mt.order))
    40  	}
    41  	*mt.stack = append(*mt.stack, mt.order)
    42  	return context.WithValue(ctx, mt, mt.order)
    43  }
    44  
    45  func (mt *mockTracer) Finish(ctx context.Context, c *app.RequestContext) {
    46  	if mt.panicAtFinish {
    47  		panic(fmt.Sprintf("panicked at finish: Tracer(%d)", mt.order))
    48  	}
    49  	*mt.stack = append(*mt.stack, -mt.order)
    50  }
    51  
    52  func TestOrder(t *testing.T) {
    53  	var c Controller
    54  	var stack []int
    55  	t1 := &mockTracer{order: 1, stack: &stack}
    56  	t2 := &mockTracer{order: 2, stack: &stack}
    57  	ctx := app.NewContext(16)
    58  	c.Append(t1)
    59  	c.Append(t2)
    60  
    61  	ctx0 := context.Background()
    62  	ctx1 := c.DoStart(ctx0, ctx)
    63  	assert.Assert(t, ctx1 != ctx0)
    64  	assert.Assert(t, len(stack) == 2 && stack[0] == 1 && stack[1] == 2, stack)
    65  
    66  	c.DoFinish(ctx1, ctx, nil)
    67  	assert.Assert(t, len(stack) == 4 && stack[2] == -2 && stack[3] == -1, stack)
    68  }
    69  
    70  func TestPanic(t *testing.T) {
    71  	var c Controller
    72  	var stack []int
    73  	t1 := &mockTracer{order: 1, stack: &stack, panicAtStart: true, panicAtFinish: true}
    74  	t2 := &mockTracer{order: 2, stack: &stack}
    75  	ctx := app.NewContext(16)
    76  	ctx.SetTraceInfo(traceinfo.NewTraceInfo())
    77  	c.Append(t1)
    78  	c.Append(t2)
    79  
    80  	ctx0 := context.Background()
    81  	ctx1 := c.DoStart(ctx0, ctx)
    82  	assert.Assert(t, ctx1 != ctx0)
    83  	assert.Assert(t, len(stack) == 0) // t1's panic skips all subsequent Starts
    84  
    85  	err := errors.New("some error")
    86  	c.DoFinish(ctx1, ctx, err)
    87  	assert.Assert(t, len(stack) == 1 && stack[0] == -2, stack)
    88  }