github.com/openfga/openfga@v1.5.4-rc1/internal/mocks/mock_tracing_server.go (about)

     1  package mocks
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"log"
     7  	"net"
     8  	"sync"
     9  	"testing"
    10  
    11  	"github.com/stretchr/testify/require"
    12  
    13  	otlpcollector "go.opentelemetry.io/proto/otlp/collector/trace/v1"
    14  	"google.golang.org/grpc"
    15  )
    16  
    17  type mockTracingServer struct {
    18  	otlpcollector.UnimplementedTraceServiceServer
    19  	exportCount int
    20  	serviceMu   sync.Mutex
    21  	server      *grpc.Server
    22  }
    23  
    24  var _ otlpcollector.TraceServiceServer = (*mockTracingServer)(nil)
    25  
    26  func (s *mockTracingServer) Export(context.Context, *otlpcollector.ExportTraceServiceRequest) (*otlpcollector.ExportTraceServiceResponse, error) {
    27  	s.serviceMu.Lock()
    28  	defer s.serviceMu.Unlock()
    29  	s.exportCount++
    30  	return &otlpcollector.ExportTraceServiceResponse{}, nil
    31  }
    32  
    33  func NewMockTracingServer(t testing.TB, port int) *mockTracingServer {
    34  	mockServer := &mockTracingServer{exportCount: 0, server: grpc.NewServer()}
    35  	otlpcollector.RegisterTraceServiceServer(mockServer.server, mockServer)
    36  	listener, err := net.Listen("tcp", fmt.Sprintf(":%d", port))
    37  	require.NoError(t, err)
    38  	t.Cleanup(mockServer.server.Stop)
    39  
    40  	go func() {
    41  		if err := mockServer.server.Serve(listener); err != nil {
    42  			log.Fatalf("failed to serve: %v", err)
    43  		}
    44  		log.Println("server closed")
    45  	}()
    46  
    47  	return mockServer
    48  }
    49  
    50  func (s *mockTracingServer) GetExportCount() int {
    51  	s.serviceMu.Lock()
    52  	defer s.serviceMu.Unlock()
    53  	return s.exportCount
    54  }