test(client-proxy): expand unit test coverage across packages

packages/client-proxy/internal/cfg/model_test.go

@@ -0,0 +1,70 @@
+package cfg
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestParse_Defaults(t *testing.T) {
+	t.Setenv("HEALTH_PORT", "")
+	t.Setenv("PROXY_PORT", "")
+	t.Setenv("REDIS_URL", "")
+	t.Setenv("REDIS_CLUSTER_URL", "")
+	t.Setenv("REDIS_TLS_CA_BASE64", "")
+	t.Setenv("REDIS_POOL_SIZE", "")
+	t.Setenv("API_INTERNAL_GRPC_ADDRESS", "")
+	t.Setenv("API_EDGE_GRPC_ADDRESS", "")
+	t.Setenv("API_EDGE_GRPC_OAUTH_CLIENT_ID", "")
+	t.Setenv("API_EDGE_GRPC_OAUTH_CLIENT_SECRET", "")
+	t.Setenv("API_EDGE_GRPC_OAUTH_TOKEN_URL", "")
+
+	cfg, err := Parse()
+	require.NoError(t, err)
+	require.EqualValues(t, 3003, cfg.HealthPort)
+	require.EqualValues(t, 3002, cfg.ProxyPort)
+	require.Equal(t, 40, cfg.RedisPoolSize)
+	require.Empty(t, cfg.RedisURL)
+	require.Empty(t, cfg.RedisClusterURL)
+	require.Empty(t, cfg.RedisTLSCABase64)
+	require.Empty(t, cfg.APIInternalGRPCAddress)
+	require.Empty(t, cfg.APIEdgeGRPCAddress)
+	require.Empty(t, cfg.APIEdgeGRPCOAuthClientID)
+	require.Empty(t, cfg.APIEdgeGRPCOAuthClientSecret)
+	require.Empty(t, cfg.APIEdgeGRPCOAuthTokenURL)
+}
+
+func TestParse_OverridesFromEnv(t *testing.T) {
+	t.Setenv("HEALTH_PORT", "9001")
+	t.Setenv("PROXY_PORT", "9002")
+	t.Setenv("REDIS_URL", "redis://localhost:6379")
+	t.Setenv("REDIS_CLUSTER_URL", "redis://cluster:6379")
+	t.Setenv("REDIS_TLS_CA_BASE64", "Y2EtZGF0YQ==")
+	t.Setenv("REDIS_POOL_SIZE", "12")
+	t.Setenv("API_INTERNAL_GRPC_ADDRESS", "internal:5005")
+	t.Setenv("API_EDGE_GRPC_ADDRESS", "edge:5006")
+	t.Setenv("API_EDGE_GRPC_OAUTH_CLIENT_ID", "client-id")
+	t.Setenv("API_EDGE_GRPC_OAUTH_CLIENT_SECRET", "client-secret")
+	t.Setenv("API_EDGE_GRPC_OAUTH_TOKEN_URL", "https://tokens.example.com")
+
+	cfg, err := Parse()
+	require.NoError(t, err)
+	require.EqualValues(t, 9001, cfg.HealthPort)
+	require.EqualValues(t, 9002, cfg.ProxyPort)
+	require.Equal(t, "redis://localhost:6379", cfg.RedisURL)
+	require.Equal(t, "redis://cluster:6379", cfg.RedisClusterURL)
+	require.Equal(t, "Y2EtZGF0YQ==", cfg.RedisTLSCABase64)
+	require.Equal(t, 12, cfg.RedisPoolSize)
+	require.Equal(t, "internal:5005", cfg.APIInternalGRPCAddress)
+	require.Equal(t, "edge:5006", cfg.APIEdgeGRPCAddress)
+	require.Equal(t, "client-id", cfg.APIEdgeGRPCOAuthClientID)
+	require.Equal(t, "client-secret", cfg.APIEdgeGRPCOAuthClientSecret)
+	require.Equal(t, "https://tokens.example.com", cfg.APIEdgeGRPCOAuthTokenURL)
+}
+
+func TestParse_InvalidIntegerReturnsError(t *testing.T) {
+	t.Setenv("HEALTH_PORT", "not-a-number")
+
+	_, err := Parse()
+	require.Error(t, err)
+}

packages/client-proxy/internal/info_test.go

@@ -0,0 +1,66 @@
+package internal
+
+import (
+	"sync"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestServiceInfo_DefaultStatusIsZeroValue(t *testing.T) {
+	t.Parallel()
+
+	s := &ServiceInfo{}
+	require.Equal(t, ServiceHealth(""), s.GetStatus())
+}
+
+func TestServiceInfo_SetAndGetStatus(t *testing.T) {
+	t.Parallel()
+
+	s := &ServiceInfo{}
+	ctx := t.Context()
+
+	s.SetStatus(ctx, Healthy)
+	require.Equal(t, Healthy, s.GetStatus())
+
+	s.SetStatus(ctx, Draining)
+	require.Equal(t, Draining, s.GetStatus())
+
+	s.SetStatus(ctx, Unhealthy)
+	require.Equal(t, Unhealthy, s.GetStatus())
+}
+
+func TestServiceInfo_SetSameStatusIsIdempotent(t *testing.T) {
+	t.Parallel()
+
+	s := &ServiceInfo{}
+	ctx := t.Context()
+
+	s.SetStatus(ctx, Healthy)
+	s.SetStatus(ctx, Healthy)
+	require.Equal(t, Healthy, s.GetStatus())
+}
+
+func TestServiceInfo_ConcurrentAccess(t *testing.T) {
+	t.Parallel()
+
+	s := &ServiceInfo{}
+	ctx := t.Context()
+	statuses := []ServiceHealth{Healthy, Draining, Unhealthy}
+
+	var wg sync.WaitGroup
+	for i := range 50 {
+		wg.Add(2)
+		go func(idx int) {
+			defer wg.Done()
+			s.SetStatus(ctx, statuses[idx%len(statuses)])
+		}(i)
+		go func() {
+			defer wg.Done()
+			_ = s.GetStatus()
+		}()
+	}
+	wg.Wait()
+
+	require.Contains(t, statuses, s.GetStatus())
+}

packages/client-proxy/internal/proxy/paused_sandbox_resumer_grpc_test.go

@@ -0,0 +1,245 @@
+package proxy
+
+import (
+	"context"
+	"errors"
+	"net"
+	"sync/atomic"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/credentials/insecure"
+	"google.golang.org/grpc/metadata"
+	"google.golang.org/grpc/status"
+	"google.golang.org/grpc/test/bufconn"
+
+	proxygrpc "github.com/e2b-dev/infra/packages/shared/pkg/grpc/proxy"
+)
+
+type fakeSandboxServer struct {
+	proxygrpc.UnimplementedSandboxServiceServer
+
+	resp *proxygrpc.SandboxResumeResponse
+	err  error
+
+	gotSandboxID string
+	gotMetadata  metadata.MD
+	calls        atomic.Int32
+}
+
+func (f *fakeSandboxServer) ResumeSandbox(ctx context.Context, req *proxygrpc.SandboxResumeRequest) (*proxygrpc.SandboxResumeResponse, error) {
+	f.calls.Add(1)
+	f.gotSandboxID = req.GetSandboxId()
+	if md, ok := metadata.FromIncomingContext(ctx); ok {
+		f.gotMetadata = md
+	}
+
+	if f.err != nil {
+		return nil, f.err
+	}
+
+	return f.resp, nil
+}
+
+func startFakeServer(t *testing.T, srv proxygrpc.SandboxServiceServer) *grpc.ClientConn {
+	t.Helper()
+
+	listener := bufconn.Listen(1024 * 1024)
+	server := grpc.NewServer()
+	proxygrpc.RegisterSandboxServiceServer(server, srv)
+
+	go func() {
+		_ = server.Serve(listener)
+	}()
+	t.Cleanup(server.Stop)
+
+	conn, err := grpc.NewClient(
+		"passthrough:///bufnet",
+		grpc.WithTransportCredentials(insecure.NewCredentials()),
+		grpc.WithContextDialer(func(ctx context.Context, _ string) (net.Conn, error) {
+			return listener.DialContext(ctx)
+		}),
+	)
+	require.NoError(t, err)
+	t.Cleanup(func() {
+		_ = conn.Close()
+	})
+
+	return conn
+}
+
+func TestNewGRPCPausedSandboxResumer_EmptyAddressErrors(t *testing.T) {
+	t.Parallel()
+
+	r, err := NewGRPCPausedSandboxResumer(t.Context(), "  ", GRPCOAuthConfig{}, false)
+	require.Error(t, err)
+	require.Nil(t, r)
+}
+
+func TestNewGRPCPausedSandboxResumer_OAuthMisconfiguredErrors(t *testing.T) {
+	t.Parallel()
+
+	r, err := NewGRPCPausedSandboxResumer(t.Context(), "127.0.0.1:1234", GRPCOAuthConfig{ClientID: "only-id"}, false)
+	require.Error(t, err)
+	require.Nil(t, r)
+}
+
+func TestNewGRPCPausedSandboxResumer_InsecureSucceeds(t *testing.T) {
+	t.Parallel()
+
+	r, err := NewGRPCPausedSandboxResumer(t.Context(), "127.0.0.1:1234", GRPCOAuthConfig{}, false)
+	require.NoError(t, err)
+	require.NotNil(t, r)
+
+	closer, ok := r.(interface {
+		Close(ctx context.Context) error
+	})
+	require.True(t, ok)
+	require.NoError(t, closer.Close(t.Context()))
+}
+
+func TestNewGRPCPausedSandboxResumer_TLSSucceeds(t *testing.T) {
+	t.Parallel()
+
+	r, err := NewGRPCPausedSandboxResumer(t.Context(), "127.0.0.1:1234", GRPCOAuthConfig{}, true)
+	require.NoError(t, err)
+	require.NotNil(t, r)
+
+	closer, ok := r.(interface {
+		Close(ctx context.Context) error
+	})
+	require.True(t, ok)
+	require.NoError(t, closer.Close(t.Context()))
+}
+
+func TestGRPCPausedSandboxResumer_ResumeSendsMetadataAndReturnsIP(t *testing.T) {
+	t.Parallel()
+
+	srv := &fakeSandboxServer{
+		resp: &proxygrpc.SandboxResumeResponse{OrchestratorIp: "  10.0.0.5  "},
+	}
+	conn := startFakeServer(t, srv)
+
+	r := &grpcPausedSandboxResumer{
+		conn:   conn,
+		client: proxygrpc.NewSandboxServiceClient(conn),
+		auth:   noopGrpcResumeAuth{},
+	}
+
+	ip, err := r.Resume(t.Context(), "sbx-123", 49983, "traffic-token", "envd-token")
+	require.NoError(t, err)
+	require.Equal(t, "10.0.0.5", ip)
+	require.EqualValues(t, 1, srv.calls.Load())
+	require.Equal(t, "sbx-123", srv.gotSandboxID)
+	require.Equal(t, []string{"49983"}, srv.gotMetadata.Get(proxygrpc.MetadataSandboxRequestPort))
+	require.Equal(t, []string{"traffic-token"}, srv.gotMetadata.Get(proxygrpc.MetadataTrafficAccessToken))
+	require.Equal(t, []string{"envd-token"}, srv.gotMetadata.Get(proxygrpc.MetadataEnvdAccessToken))
+}
+
+func TestGRPCPausedSandboxResumer_ResumeOmitsEmptyTokens(t *testing.T) {
+	t.Parallel()
+
+	srv := &fakeSandboxServer{
+		resp: &proxygrpc.SandboxResumeResponse{OrchestratorIp: "10.0.0.6"},
+	}
+	conn := startFakeServer(t, srv)
+
+	r := &grpcPausedSandboxResumer{
+		conn:   conn,
+		client: proxygrpc.NewSandboxServiceClient(conn),
+		auth:   noopGrpcResumeAuth{},
+	}
+
+	ip, err := r.Resume(t.Context(), "sbx-456", 8080, "", "")
+	require.NoError(t, err)
+	require.Equal(t, "10.0.0.6", ip)
+	require.EqualValues(t, 1, srv.calls.Load())
+	require.Empty(t, srv.gotMetadata.Get(proxygrpc.MetadataTrafficAccessToken))
+	require.Empty(t, srv.gotMetadata.Get(proxygrpc.MetadataEnvdAccessToken))
+	require.Equal(t, []string{"8080"}, srv.gotMetadata.Get(proxygrpc.MetadataSandboxRequestPort))
+}
+
+func TestGRPCPausedSandboxResumer_ResumeReturnsServerError(t *testing.T) {
+	t.Parallel()
+
+	srv := &fakeSandboxServer{
+		err: status.Error(codes.NotFound, "missing"),
+	}
+	conn := startFakeServer(t, srv)
+
+	r := &grpcPausedSandboxResumer{
+		conn:   conn,
+		client: proxygrpc.NewSandboxServiceClient(conn),
+		auth:   noopGrpcResumeAuth{},
+	}
+
+	ip, err := r.Resume(t.Context(), "sbx", 80, "", "")
+	require.Error(t, err)
+	require.Empty(t, ip)
+	st, ok := status.FromError(errors.Unwrap(err))
+	require.True(t, ok)
+	require.Equal(t, codes.NotFound, st.Code())
+	require.EqualValues(t, 1, srv.calls.Load())
+}
+
+type errAuth struct {
+	err error
+}
+
+func (e errAuth) authorize(_ context.Context) (context.Context, error) {
+	return nil, e.err
+}
+
+func TestGRPCPausedSandboxResumer_ResumeAuthError(t *testing.T) {
+	t.Parallel()
+
+	srv := &fakeSandboxServer{
+		resp: &proxygrpc.SandboxResumeResponse{OrchestratorIp: "10.0.0.7"},
+	}
+	conn := startFakeServer(t, srv)
+
+	authErr := errors.New("auth failed")
+	r := &grpcPausedSandboxResumer{
+		conn:   conn,
+		client: proxygrpc.NewSandboxServiceClient(conn),
+		auth:   errAuth{err: authErr},
+	}
+
+	ip, err := r.Resume(t.Context(), "sbx", 80, "", "")
+	require.ErrorIs(t, err, authErr)
+	require.Empty(t, ip)
+	require.EqualValues(t, 0, srv.calls.Load())
+}
+
+func TestGRPCPausedSandboxResumer_Init(t *testing.T) {
+	t.Parallel()
+
+	srv := &fakeSandboxServer{}
+	conn := startFakeServer(t, srv)
+
+	r := &grpcPausedSandboxResumer{
+		conn:   conn,
+		client: proxygrpc.NewSandboxServiceClient(conn),
+		auth:   noopGrpcResumeAuth{},
+	}
+
+	// Should not panic
+	r.Init(t.Context())
+}
+
+func TestGRPCPausedSandboxResumer_Close(t *testing.T) {
+	t.Parallel()
+
+	srv := &fakeSandboxServer{}
+	conn := startFakeServer(t, srv)
+
+	r := &grpcPausedSandboxResumer{
+		conn:   conn,
+		client: proxygrpc.NewSandboxServiceClient(conn),
+		auth:   noopGrpcResumeAuth{},
+	}
+
+	require.NoError(t, r.Close(t.Context()))
+}