github.com/castai/kvisor@v1.7.1-0.20240516114728-b3572a2607b5/cmd/agent/daemon/state/clickhouse_netflow_exporter_test.go (about)

     1  package state
     2  
     3  import (
     4  	"context"
     5  	"net/netip"
     6  	"testing"
     7  	"time"
     8  
     9  	"github.com/ClickHouse/clickhouse-go/v2"
    10  	castpb "github.com/castai/kvisor/api/v1/runtime"
    11  	"github.com/castai/kvisor/pkg/logging"
    12  	"github.com/stretchr/testify/require"
    13  	"github.com/testcontainers/testcontainers-go"
    14  	"github.com/testcontainers/testcontainers-go/wait"
    15  )
    16  
    17  func TestClickhouseNetflowExporter(t *testing.T) {
    18  	if testing.Short() {
    19  		t.Skip()
    20  	}
    21  
    22  	r := require.New(t)
    23  	ctx := context.Background()
    24  	log := logging.NewTestLog()
    25  
    26  	addr, cleanup := startClickhouseDB(t)
    27  	defer cleanup()
    28  
    29  	conn, err := newClickhouseConn(addr)
    30  	r.NoError(err)
    31  
    32  	r.NoError(conn.Exec(ctx, ClickhouseNetflowSchema()))
    33  
    34  	exporter := NewClickhouseNetflowExporter(log, conn, 100)
    35  	err = exporter.asyncWrite(ctx, true, &castpb.Netflow{
    36  		StartTs:       uint64(time.Now().UnixNano()),
    37  		EndTs:         uint64(time.Now().Add(time.Minute).UnixNano()),
    38  		ProcessName:   "curl",
    39  		Namespace:     "n1",
    40  		PodName:       "p1",
    41  		ContainerName: "c1",
    42  		WorkloadName:  "w1",
    43  		WorkloadKind:  "Deployment",
    44  		Zone:          "us-east-1",
    45  		Addr:          netip.MustParseAddr("10.10.1.10").AsSlice(),
    46  		Port:          345555,
    47  		Protocol:      castpb.NetflowProtocol_NETFLOW_PROTOCOL_TCP,
    48  		Destinations: []*castpb.NetflowDestination{
    49  			{
    50  				Namespace:    "n2",
    51  				PodName:      "p2",
    52  				WorkloadName: "w2",
    53  				WorkloadKind: "Deployment",
    54  				Zone:         "us-east-2",
    55  				DnsQuestion:  "service.n2.svc.cluster.local.",
    56  				Addr:         netip.MustParseAddr("10.10.1.15").AsSlice(),
    57  				Port:         80,
    58  				TxBytes:      10,
    59  				TxPackets:    2,
    60  			},
    61  		},
    62  	})
    63  	r.NoError(err)
    64  }
    65  
    66  func newClickhouseConn(clickhouseAddr string) (clickhouse.Conn, error) {
    67  	return clickhouse.Open(&clickhouse.Options{
    68  		Addr: []string{clickhouseAddr},
    69  		Auth: clickhouse.Auth{
    70  			Database: "kvisor",
    71  			Username: "kvisor",
    72  			Password: "kvisor",
    73  		},
    74  		Settings: clickhouse.Settings{
    75  			"allow_experimental_object_type": "1",
    76  		},
    77  		MaxOpenConns: 20,
    78  	})
    79  }
    80  
    81  func startClickhouseDB(t *testing.T) (string, func()) {
    82  	ctx := context.Background()
    83  	hz := wait.NewHTTPStrategy("/ping")
    84  	hz.Port = "8123/tcp"
    85  	req := testcontainers.ContainerRequest{
    86  		Image:        "clickhouse/clickhouse-server:24.2.3.70-alpine",
    87  		ExposedPorts: []string{"9000/tcp", "8123/tcp"},
    88  		WaitingFor:   hz,
    89  		Env: map[string]string{
    90  			"CLICKHOUSE_USER":                      "kvisor",
    91  			"CLICKHOUSE_PASSWORD":                  "kvisor",
    92  			"CLICKHOUSE_DB":                        "kvisor",
    93  			"CLICKHOUSE_DEFAULT_ACCESS_MANAGEMENT": "1",
    94  		},
    95  	}
    96  	cont, err := testcontainers.GenericContainer(ctx, testcontainers.GenericContainerRequest{
    97  		ContainerRequest: req,
    98  		Started:          true,
    99  	})
   100  	if err != nil {
   101  		t.Fatal(err)
   102  	}
   103  	mport, err := cont.MappedPort(ctx, "9000/tcp")
   104  	if err != nil {
   105  		t.Fatal(err)
   106  	}
   107  	return "127.0.0.1:" + mport.Port(), func() {
   108  		if err := cont.Terminate(ctx); err != nil {
   109  			t.Fatalf("failed to terminate container: %s", err.Error())
   110  		}
   111  	}
   112  }