Refactor API & FLAPS clients to interfaces for unit testing

agent/client.go

@@ -18,12 +18,12 @@ import (
 	"github.com/azazeal/pause"
 	"golang.org/x/sync/errgroup"
 
-	fly "github.com/superfly/fly-go"
 	"github.com/superfly/flyctl/agent/internal/proto"
 	"github.com/superfly/flyctl/gql"
 	"github.com/superfly/flyctl/internal/buildinfo"
 	"github.com/superfly/flyctl/internal/config"
 	"github.com/superfly/flyctl/internal/flag"
+	"github.com/superfly/flyctl/internal/flyutil"
 	"github.com/superfly/flyctl/internal/logger"
 	"github.com/superfly/flyctl/internal/sentry"
 	"github.com/superfly/flyctl/internal/version"
@@ -34,7 +34,7 @@ import (
 )
 
 // Establish starts the daemon, if necessary, and returns a client to it.
-func Establish(ctx context.Context, apiClient *fly.Client) (*Client, error) {
+func Establish(ctx context.Context, apiClient flyutil.Client) (*Client, error) {
 	if err := wireguard.PruneInvalidPeers(ctx, apiClient); err != nil {
 		return nil, err
 	}
@@ -540,7 +540,7 @@ func arrayEqual(a, b []string) bool {
 }
 
 func gqlGetInstances(ctx context.Context, orgSlug, appName string) instancesResult {
-	gqlClient := fly.ClientFromContext(ctx).GenqClient
+	gqlClient := flyutil.ClientFromContext(ctx).GenqClient()
 	_ = `# @genqlient
 	query AgentGetInstances($appName: String!) {
 		app(name: $appName) {

agent/server/server.go

@@ -224,7 +224,7 @@ func (s *server) checkForConfigChange() (err error) {
 	return
 }
 
-func (s *server) buildTunnel(ctx context.Context, org *fly.Organization, recycle bool, network string, client *fly.Client) (tunnel *wg.Tunnel, err error) {
+func (s *server) buildTunnel(ctx context.Context, org *fly.Organization, recycle bool, network string, client flyutil.Client) (tunnel *wg.Tunnel, err error) {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 
@@ -392,7 +392,7 @@ func (s *server) clean(ctx context.Context) {
 
 // GetClient returns an API client that uses the server's tokens. Sessions may
 // have their own tokens, so should use session.getClient instead.
-func (s *server) GetClient(ctx context.Context) *fly.Client {
+func (s *server) GetClient(ctx context.Context) flyutil.Client {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 

agent/server/session.go

@@ -579,7 +579,7 @@ func (s *session) setToken(ctx context.Context, args ...string) {
 
 // getClient returns an API client that uses any API tokens sent by the client.
 // If none have been sent, it falls back to using the server's tokens.
-func (s *session) getClient(ctx context.Context) *fly.Client {
+func (s *session) getClient(ctx context.Context) flyutil.Client {
 	if s.tokens == nil {
 		return s.srv.GetClient(ctx)
 	}

flypg/cmd.go

@@ -10,6 +10,7 @@ import (
 	"github.com/superfly/flyctl/agent"
 	"github.com/superfly/flyctl/internal/command/ssh"
 	"github.com/superfly/flyctl/internal/flag"
+	"github.com/superfly/flyctl/internal/flyutil"
 	"github.com/superfly/flyctl/iostreams"
 )
 
@@ -31,7 +32,7 @@ type Command struct {
 }
 
 func NewCommand(ctx context.Context, app *fly.AppCompact) (*Command, error) {
-	client := fly.ClientFromContext(ctx)
+	client := flyutil.ClientFromContext(ctx)
 
 	agentclient, err := agent.Establish(ctx, client)
 	if err != nil {

flypg/launcher.go

@@ -12,6 +12,7 @@ import (
 	"github.com/superfly/flyctl/helpers"
 	"github.com/superfly/flyctl/internal/buildinfo"
 	"github.com/superfly/flyctl/internal/flapsutil"
+	"github.com/superfly/flyctl/internal/flyutil"
 	mach "github.com/superfly/flyctl/internal/machine"
 	"github.com/superfly/flyctl/internal/watch"
 
@@ -35,7 +36,7 @@ const (
 )
 
 type Launcher struct {
-	client *fly.Client
+	client flyutil.Client
 }
 
 type CreateClusterInput struct {
@@ -55,7 +56,7 @@ type CreateClusterInput struct {
 	ForkFrom           string
 }
 
-func NewLauncher(client *fly.Client) *Launcher {
+func NewLauncher(client flyutil.Client) *Launcher {
 	return &Launcher{
 		client: client,
 	}
@@ -66,7 +67,7 @@ func (l *Launcher) LaunchMachinesPostgres(ctx context.Context, config *CreateClu
 	var (
 		io       = iostreams.FromContext(ctx)
 		colorize = io.ColorScheme()
-		client   = fly.ClientFromContext(ctx)
+		client   = flyutil.ClientFromContext(ctx)
 	)
 
 	// Ensure machines can be started when scaling to zero is enabled
@@ -102,7 +103,7 @@ func (l *Launcher) LaunchMachinesPostgres(ctx context.Context, config *CreateClu
 	if err != nil {
 		return err
 	}
-	ctx = flaps.NewContext(ctx, flapsClient)
+	ctx = flapsutil.NewContextWithClient(ctx, flapsClient)
 
 	nodes := make([]*fly.Machine, 0)
 

go.mod

@@ -63,7 +63,7 @@ require (
 	github.com/spf13/pflag v1.0.5
 	github.com/spf13/viper v1.18.2
 	github.com/stretchr/testify v1.9.0
-	github.com/superfly/fly-go v0.1.11
+	github.com/superfly/fly-go v0.1.12
 	github.com/superfly/graphql v0.2.4
 	github.com/superfly/lfsc-go v0.1.1
 	github.com/superfly/macaroon v0.2.13

go.sum

@@ -614,8 +614,8 @@ github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsT
 github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
 github.com/subosito/gotenv v1.6.0 h1:9NlTDc1FTs4qu0DDq7AEtTPNw6SVm7uBMsUCUjABIf8=
 github.com/subosito/gotenv v1.6.0/go.mod h1:Dk4QP5c2W3ibzajGcXpNraDfq2IrhjMIvMSWPKKo0FU=
-github.com/superfly/fly-go v0.1.11 h1:gYoTIHqsaLvmfQWmKwX4uGWdzvu+hUnklYtwX+sUtJE=
-github.com/superfly/fly-go v0.1.11/go.mod h1:0wgmVsqPPDgPSTsZA0L0zI5uajxx86KfuBoPQT87B1E=
+github.com/superfly/fly-go v0.1.12 h1:Ql+wbDFuiWEGK+qskuSTtMft/zfh1rIBp/1Dcu7JbgQ=
+github.com/superfly/fly-go v0.1.12/go.mod h1:0wgmVsqPPDgPSTsZA0L0zI5uajxx86KfuBoPQT87B1E=
 github.com/superfly/graphql v0.2.4 h1:Av8hSk4x8WvKJ6MTnEwrLknSVSGPc7DWpgT3z/kt3PU=
 github.com/superfly/graphql v0.2.4/go.mod h1:CVfDl31srm8HnJ9udwLu6hFNUW/P6GUM2dKcG1YQ8jc=
 github.com/superfly/lfsc-go v0.1.1 h1:dGjLgt81D09cG+aR9lJZIdmonjZSR5zYCi7s54+ZU2Q=

internal/appconfig/remote.go

@@ -5,14 +5,15 @@ import (
 	"fmt"
 
 	fly "github.com/superfly/fly-go"
-	"github.com/superfly/fly-go/flaps"
 	"github.com/superfly/flyctl/gql"
+	"github.com/superfly/flyctl/internal/flapsutil"
+	"github.com/superfly/flyctl/internal/flyutil"
 	"github.com/superfly/flyctl/internal/machine"
 	"github.com/superfly/flyctl/iostreams"
 )
 
 func FromRemoteApp(ctx context.Context, appName string) (*Config, error) {
-	apiClient := fly.ClientFromContext(ctx)
+	apiClient := flyutil.ClientFromContext(ctx)
 
 	cfg, err := getAppV2ConfigFromReleases(ctx, apiClient, appName)
 	if cfg == nil {
@@ -29,7 +30,7 @@ func FromRemoteApp(ctx context.Context, appName string) (*Config, error) {
 }
 
 func getAppV2ConfigFromMachines(ctx context.Context, appName string) (*Config, error) {
-	flapsClient := flaps.FromContext(ctx)
+	flapsClient := flapsutil.ClientFromContext(ctx)
 	io := iostreams.FromContext(ctx)
 
 	activeMachines, err := machine.ListActive(ctx)
@@ -47,7 +48,7 @@ func getAppV2ConfigFromMachines(ctx context.Context, appName string) (*Config, e
 	return appConfig, nil
 }
 
-func getAppV2ConfigFromReleases(ctx context.Context, apiClient *fly.Client, appName string) (*Config, error) {
+func getAppV2ConfigFromReleases(ctx context.Context, apiClient flyutil.Client, appName string) (*Config, error) {
 	_ = `# @genqlient
 	query FlyctlConfigCurrentRelease($appName: String!) {
 		app(name:$appName) {
@@ -57,7 +58,7 @@ func getAppV2ConfigFromReleases(ctx context.Context, apiClient *fly.Client, appN
 		}
 	}
 	`
-	resp, err := gql.FlyctlConfigCurrentRelease(ctx, apiClient.GenqClient, appName)
+	resp, err := gql.FlyctlConfigCurrentRelease(ctx, apiClient.GenqClient(), appName)
 	if err != nil {
 		return nil, err
 	}

internal/build/imgsrc/docker.go

@@ -30,6 +30,7 @@ import (
 	"github.com/superfly/flyctl/helpers"
 	"github.com/superfly/flyctl/internal/config"
 	"github.com/superfly/flyctl/internal/flyerr"
+	"github.com/superfly/flyctl/internal/flyutil"
 	"github.com/superfly/flyctl/internal/metrics"
 	"github.com/superfly/flyctl/internal/sentry"
 	"github.com/superfly/flyctl/internal/tracing"
@@ -48,11 +49,11 @@ type dockerClientFactory struct {
 	mode      DockerDaemonType
 	remote    bool
 	buildFn   func(ctx context.Context, build *build) (*dockerclient.Client, error)
-	apiClient *fly.Client
+	apiClient flyutil.Client
 	appName   string
 }
 
-func newDockerClientFactory(daemonType DockerDaemonType, apiClient *fly.Client, appName string, streams *iostreams.IOStreams, connectOverWireguard bool) *dockerClientFactory {
+func newDockerClientFactory(daemonType DockerDaemonType, apiClient flyutil.Client, appName string, streams *iostreams.IOStreams, connectOverWireguard bool) *dockerClientFactory {
 	remoteFactory := func() *dockerClientFactory {
 		terminal.Debug("trying remote docker daemon")
 		return &dockerClientFactory{
@@ -202,7 +203,7 @@ func logClearLinesAbove(streams *iostreams.IOStreams, count int) {
 	}
 }
 
-func newRemoteDockerClient(ctx context.Context, apiClient *fly.Client, appName string, streams *iostreams.IOStreams, build *build, cachedClient *dockerclient.Client, connectOverWireguard bool) (c *dockerclient.Client, err error) {
+func newRemoteDockerClient(ctx context.Context, apiClient flyutil.Client, appName string, streams *iostreams.IOStreams, build *build, cachedClient *dockerclient.Client, connectOverWireguard bool) (c *dockerclient.Client, err error) {
 	ctx, span := tracing.GetTracer().Start(ctx, "build_remote_docker_client", trace.WithAttributes(
 		attribute.Bool("connect_over_wireguard", connectOverWireguard),
 	))
@@ -476,10 +477,9 @@ func buildWireguardlessClientOpts(ctx context.Context, host, appName string) ([]
 	}
 
 	return opts, nil
-
 }
 
-func buildRemoteClientOpts(ctx context.Context, apiClient *fly.Client, appName, host string) (opts []dockerclient.Opt, err error) {
+func buildRemoteClientOpts(ctx context.Context, apiClient flyutil.Client, appName, host string) (opts []dockerclient.Opt, err error) {
 	ctx, span := tracing.GetTracer().Start(ctx, "build_remote_client_ops")
 	defer span.End()
 
@@ -698,7 +698,7 @@ func ResolveDockerfile(cwd string) string {
 	return ""
 }
 
-func EagerlyEnsureRemoteBuilder(ctx context.Context, apiClient *fly.Client, orgSlug string) {
+func EagerlyEnsureRemoteBuilder(ctx context.Context, apiClient flyutil.Client, orgSlug string) {
 	// skip if local docker is available
 	if _, err := NewLocalDockerClient(); err == nil {
 		return
@@ -720,7 +720,7 @@ func EagerlyEnsureRemoteBuilder(ctx context.Context, apiClient *fly.Client, orgS
 	terminal.Debugf("remote builder %s is being prepared", app.Name)
 }
 
-func remoteBuilderMachine(ctx context.Context, apiClient *fly.Client, appName string) (*fly.GqlMachine, *fly.App, error) {
+func remoteBuilderMachine(ctx context.Context, apiClient flyutil.Client, appName string) (*fly.GqlMachine, *fly.App, error) {
 	if v := os.Getenv("FLY_REMOTE_BUILDER_HOST"); v != "" {
 		return nil, nil, nil
 	}

internal/build/imgsrc/remote_image_resolver.go

@@ -5,14 +5,14 @@ import (
 	"fmt"
 	"strconv"
 
-	fly "github.com/superfly/fly-go"
+	"github.com/superfly/flyctl/internal/flyutil"
 	"github.com/superfly/flyctl/internal/tracing"
 	"github.com/superfly/flyctl/iostreams"
 	"go.opentelemetry.io/otel/trace"
 )
 
 type remoteImageResolver struct {
-	flyApi *fly.Client
+	flyApi flyutil.Client
 }
 
 func (*remoteImageResolver) Name() string {

internal/build/imgsrc/resolver.go

@@ -22,11 +22,11 @@ import (
 	"github.com/superfly/flyctl/gql"
 	"github.com/superfly/flyctl/internal/buildinfo"
 	"github.com/superfly/flyctl/internal/config"
+	"github.com/superfly/flyctl/internal/flyutil"
 	"github.com/superfly/flyctl/internal/sentry"
 	"github.com/superfly/flyctl/internal/tracing"
 	"github.com/superfly/flyctl/iostreams"
 
-	fly "github.com/superfly/fly-go"
 	"github.com/superfly/flyctl/terminal"
 )
 
@@ -144,7 +144,7 @@ func (di DeploymentImage) ToSpanAttributes() []attribute.KeyValue {
 
 type Resolver struct {
 	dockerFactory *dockerClientFactory
-	apiClient     *fly.Client
+	apiClient     flyutil.Client
 }
 
 type StopSignal struct {
@@ -310,7 +310,7 @@ func (r *Resolver) createBuildGql(ctx context.Context, strategiesAvailable []str
 	ctx, span := tracing.GetTracer().Start(ctx, "web.create_build")
 	defer span.End()
 
-	gqlClient := fly.ClientFromContext(ctx).GenqClient
+	gqlClient := flyutil.ClientFromContext(ctx).GenqClient()
 	_ = `# @genqlient
 	mutation ResolverCreateBuild($input:CreateBuildInput!) {
 		createBuild(input:$input) {
@@ -513,7 +513,7 @@ func (r *Resolver) finishBuild(ctx context.Context, build *build, failed bool, l
 		terminal.Debug("Skipping FinishBuild() gql call, because CreateBuild() failed.\n")
 		return nil, nil
 	}
-	gqlClient := fly.ClientFromContext(ctx).GenqClient
+	gqlClient := flyutil.ClientFromContext(ctx).GenqClient()
 	_ = `# @genqlient
 	mutation ResolverFinishBuild($input:FinishBuildInput!) {
 		finishBuild(input:$input) {
@@ -750,7 +750,7 @@ func (s *StopSignal) Stop() {
 	})
 }
 
-func NewResolver(daemonType DockerDaemonType, apiClient *fly.Client, appName string, iostreams *iostreams.IOStreams, connectOverWireguard bool) *Resolver {
+func NewResolver(daemonType DockerDaemonType, apiClient flyutil.Client, appName string, iostreams *iostreams.IOStreams, connectOverWireguard bool) *Resolver {
 	return &Resolver{
 		dockerFactory: newDockerClientFactory(daemonType, apiClient, appName, iostreams, connectOverWireguard),
 		apiClient:     apiClient,

internal/cmdutil/preparers/preparers.go

@@ -50,10 +50,13 @@ func InitClient(ctx context.Context) (context.Context, error) {
 	fly.SetInstrumenter(instrument.ApiAdapter)
 	fly.SetTransport(otelhttp.NewTransport(http.DefaultTransport))
 
-	c := flyutil.NewClientFromOptions(ctx, fly.ClientOptions{Tokens: cfg.Tokens})
-	logger.Debug("client initialized.")
+	if flyutil.ClientFromContext(ctx) == nil {
+		client := flyutil.NewClientFromOptions(ctx, fly.ClientOptions{Tokens: cfg.Tokens})
+		logger.Debug("client initialized.")
+		ctx = flyutil.NewContextWithClient(ctx, client)
+	}
 
-	return fly.NewContextWithClient(ctx, c), nil
+	return ctx, nil
 }
 
 func DetermineConfigDir(ctx context.Context) (context.Context, error) {

internal/command/agent/agent.go

@@ -8,11 +8,11 @@ import (
 
 	"github.com/spf13/cobra"
 
-	fly "github.com/superfly/fly-go"
 	"github.com/superfly/flyctl/agent"
 
 	"github.com/superfly/flyctl/internal/command"
 	"github.com/superfly/flyctl/internal/env"
+	"github.com/superfly/flyctl/internal/flyutil"
 	"github.com/superfly/flyctl/internal/state"
 )
 
@@ -48,7 +48,7 @@ func New() (cmd *cobra.Command) {
 }
 
 func establish(ctx context.Context) (ac *agent.Client, err error) {
-	client := fly.ClientFromContext(ctx)
+	client := flyutil.ClientFromContext(ctx)
 
 	if ac, err = agent.Establish(ctx, client); err != nil {
 		err = fmt.Errorf("failed establishing connection to agent: %w", err)

internal/command/agent/instances.go

@@ -12,6 +12,7 @@ import (
 
 	"github.com/superfly/flyctl/internal/command"
 	"github.com/superfly/flyctl/internal/flag"
+	"github.com/superfly/flyctl/internal/flyutil"
 	"github.com/superfly/flyctl/internal/render"
 )
 
@@ -38,7 +39,7 @@ func runInstances(ctx context.Context) (err error) {
 	}
 
 	slug := flag.FirstArg(ctx)
-	apiClient := fly.ClientFromContext(ctx)
+	apiClient := flyutil.ClientFromContext(ctx)
 
 	var org *fly.Organization
 	if org, err = apiClient.GetOrganizationBySlug(ctx, slug); err != nil {

internal/command/apps/apps.go

@@ -12,6 +12,7 @@ import (
 	"github.com/superfly/flyctl/agent"
 	"github.com/superfly/flyctl/internal/command"
 	"github.com/superfly/flyctl/internal/flapsutil"
+	"github.com/superfly/flyctl/internal/flyutil"
 )
 
 // New initializes and returns a new apps Command.
@@ -48,7 +49,7 @@ The LIST command will list all currently registered applications.
 
 // BuildContext is a helper that builds out commonly required context requirements
 func BuildContext(ctx context.Context, app *fly.AppCompact) (context.Context, error) {
-	client := fly.ClientFromContext(ctx)
+	client := flyutil.ClientFromContext(ctx)
 
 	agentclient, err := agent.Establish(ctx, client)
 	if err != nil {
@@ -69,7 +70,7 @@ func BuildContext(ctx context.Context, app *fly.AppCompact) (context.Context, er
 		return nil, err
 	}
 
-	ctx = flaps.NewContext(ctx, flapsClient)
+	ctx = flapsutil.NewContextWithClient(ctx, flapsClient)
 
 	return ctx, nil
 }