github.com/anycable/anycable-go@v1.5.1/telemetry/telemetry_test.go (about)

     1  package telemetry
     2  
     3  import (
     4  	"context"
     5  	"log/slog"
     6  	"testing"
     7  	"time"
     8  
     9  	"github.com/anycable/anycable-go/config"
    10  	"github.com/anycable/anycable-go/metrics"
    11  	"github.com/anycable/anycable-go/version"
    12  	"github.com/posthog/posthog-go"
    13  	"github.com/stretchr/testify/assert"
    14  	"github.com/stretchr/testify/require"
    15  )
    16  
    17  type MockClient struct {
    18  	closed   bool
    19  	captured []posthog.Capture
    20  }
    21  
    22  func NewMockClient() *MockClient {
    23  	return &MockClient{captured: []posthog.Capture{}}
    24  }
    25  
    26  func (c *MockClient) Enqueue(msg posthog.Message) error {
    27  	c.captured = append(c.captured, msg.(posthog.Capture))
    28  
    29  	return nil
    30  }
    31  
    32  func (c *MockClient) Close() error {
    33  	c.closed = true
    34  
    35  	return nil
    36  }
    37  
    38  func (c *MockClient) GetAllFlags(flags posthog.FeatureFlagPayloadNoKey) (map[string]interface{}, error) {
    39  	return nil, nil
    40  }
    41  
    42  func (c *MockClient) GetFeatureFlag(flag posthog.FeatureFlagPayload) (interface{}, error) {
    43  	return nil, nil
    44  }
    45  
    46  func (c *MockClient) GetFeatureFlags() ([]posthog.FeatureFlag, error) {
    47  	return nil, nil
    48  }
    49  
    50  func (c *MockClient) IsFeatureEnabled(flag posthog.FeatureFlagPayload) (interface{}, error) {
    51  	return nil, nil
    52  }
    53  
    54  func (c *MockClient) ReloadFeatureFlags() error {
    55  	return nil
    56  }
    57  
    58  func TestTracking(t *testing.T) {
    59  	mconfig := metrics.NewConfig()
    60  	metrics, _ := metrics.NewFromConfig(&mconfig, slog.Default())
    61  
    62  	metrics.RegisterGauge("clients_num", "")
    63  	metrics.RegisterGauge("mem_sys_bytes", "")
    64  	metrics.GaugeSet("clients_num", 10)
    65  	metrics.GaugeSet("mem_sys_bytes", 100)
    66  
    67  	t.Setenv("AWS_EXECUTION_ENV", "AWS_ECS_FARGATE")
    68  
    69  	conf := config.NewConfig()
    70  	tracker := NewTracker(metrics, &conf, &Config{})
    71  	defer tracker.Shutdown(context.Background()) // nolint: errcheck
    72  
    73  	client := NewMockClient()
    74  	tracker.client = client
    75  
    76  	tracker.Collect()
    77  
    78  	require.Equal(t, 1, len(client.captured))
    79  
    80  	event := client.captured[0]
    81  
    82  	assert.Equal(t, "boot", event.Event)
    83  	assert.Equal(t, version.Version(), event.Properties["version"])
    84  
    85  	time.Sleep(100 * time.Millisecond)
    86  
    87  	metrics.GaugeSet("clients_num", 14)
    88  	metrics.GaugeSet("mem_sys_bytes", 90)
    89  
    90  	tracker.observeUsage()
    91  	tracker.collectUsage()
    92  
    93  	require.Equal(t, 2, len(client.captured))
    94  
    95  	event = client.captured[1]
    96  
    97  	assert.Equal(t, "usage", event.Event)
    98  	assert.Equal(t, "ecs-fargate", event.Properties["deploy"])
    99  	assert.Equal(t, 14, int(event.Properties["clients_max"].(uint64)))
   100  	assert.Equal(t, 100, int(event.Properties["mem_sys_max"].(uint64)))
   101  
   102  	require.NoError(t, tracker.Shutdown(context.Background()))
   103  
   104  	assert.True(t, client.closed)
   105  }