test: test deployMachinesApp

This test is a by-product of my investigation. The investigation resulted a server-side fix, but the test itself and some doc comments may be useful.
superfly · Nov 22, 2024 · f36e8d0 · f36e8d0
1 parent 0a5966d
commit f36e8d0
Show file tree

Hide file tree

Showing 11 changed files with 142 additions and 20 deletions.
diff --git a/agent/client.go b/agent/client.go
@@ -34,7 +34,7 @@ import (
 )
 
 // Establish starts the daemon, if necessary, and returns a client to it.
-func Establish(ctx context.Context, apiClient flyutil.Client) (*Client, error) {
+func Establish(ctx context.Context, apiClient wireguard.WebClient) (*Client, error) {
 	if err := wireguard.PruneInvalidPeers(ctx, apiClient); err != nil {
 		return nil, err
 	}

diff --git a/internal/command/deploy/machines.go b/internal/command/deploy/machines.go
@@ -113,13 +113,16 @@ func argsFromManifest(manifest *DeployManifest, app *fly.AppCompact) MachineDepl
 }
 
 type machineDeployment struct {
-	apiClient             flyutil.Client
-	flapsClient           flapsutil.FlapsClient
-	io                    *iostreams.IOStreams
-	colorize              *iostreams.ColorScheme
-	app                   *fly.AppCompact
-	appConfig             *appconfig.Config
-	img                   string
+	// apiClient is a client to use web.
+	apiClient webClient
+	// flapsClient is a client to use flaps.
+	flapsClient flapsutil.FlapsClient
+	io          *iostreams.IOStreams
+	colorize    *iostreams.ColorScheme
+	app         *fly.AppCompact
+	appConfig   *appconfig.Config
+	img         string
+	// machineSet is this application's machines.
 	machineSet            machine.MachineSet
 	releaseCommandMachine machine.MachineSet
 	volumes               map[string][]fly.Volume

diff --git a/internal/command/deploy/machines_deploymachinesapp.go b/internal/command/deploy/machines_deploymachinesapp.go
@@ -1106,6 +1106,7 @@ func (md *machineDeployment) spawnMachineInGroup(ctx context.Context, groupName
 	return lm, nil
 }
 
+// resolveProcessGroupChanges returns a diff between machines
 func (md *machineDeployment) resolveProcessGroupChanges() ProcessGroupsDiff {
 	output := ProcessGroupsDiff{
 		groupsToRemove:        map[string]int{},

diff --git a/internal/command/deploy/machines_deploymachinesapp_test.go b/internal/command/deploy/machines_deploymachinesapp_test.go
@@ -3,10 +3,14 @@ package deploy
 import (
 	"context"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/superfly/fly-go"
+	"github.com/superfly/flyctl/internal/appconfig"
+	"github.com/superfly/flyctl/internal/flapsutil"
 	"github.com/superfly/flyctl/internal/machine"
+	"github.com/superfly/flyctl/internal/mock"
 	"github.com/superfly/flyctl/iostreams"
 )
 
@@ -34,3 +38,37 @@ func TestUpdateExistingMachinesWRecovery(t *testing.T) {
 	})
 	assert.Error(t, err, "failed to find machine test-machine-id")
 }
+
+func TestDeployMachinesApp(t *testing.T) {
+	ios, _, _, _ := iostreams.Test()
+	client := &mockFlapsClient{}
+	webClient := &mock.Client{
+		GetAppLogsFunc: func(ctx context.Context, appName, token, region, instanceID string) (entries []fly.LogEntry, nextToken string, err error) {
+			return nil, "", nil
+		},
+	}
+	client.machines = []*fly.Machine{
+		{ID: "m1", Config: &fly.MachineConfig{Metadata: map[string]string{fly.MachineConfigMetadataKeyFlyProcessGroup: "app"}}},
+		{ID: "m2", Config: &fly.MachineConfig{Metadata: map[string]string{fly.MachineConfigMetadataKeyFlyProcessGroup: "app"}}},
+		{ID: "m3", Config: &fly.MachineConfig{Metadata: map[string]string{fly.MachineConfigMetadataKeyFlyProcessGroup: "app"}}},
+		{ID: "m4", Config: &fly.MachineConfig{Metadata: map[string]string{fly.MachineConfigMetadataKeyFlyProcessGroup: "app"}}},
+	}
+	md := &machineDeployment{
+		app:             &fly.AppCompact{},
+		io:              ios,
+		colorize:        ios.ColorScheme(),
+		flapsClient:     client,
+		apiClient:       webClient,
+		strategy:        "canary",
+		appConfig:       &appconfig.Config{},
+		machineSet:      machine.NewMachineSet(client, ios, client.machines, false),
+		skipSmokeChecks: true,
+		waitTimeout:     1 * time.Second,
+	}
+
+	ctx := context.Background()
+	ctx = iostreams.NewContext(ctx, ios)
+	ctx = flapsutil.NewContextWithClient(ctx, client)
+	err := md.deployMachinesApp(ctx)
+	assert.NoError(t, err)
+}
diff --git a/internal/command/deploy/mock_client_test.go b/internal/command/deploy/mock_client_test.go
@@ -16,18 +16,37 @@ func (f *mockWebClient) CanPerformBluegreenDeployment(ctx context.Context, appNa
 	return true, nil
 }
 
+var (
+	globalMachineID = 0
+)
+
 type mockFlapsClient struct {
 	breakLaunch      bool
 	breakWait        bool
 	breakUncordon    bool
 	breakSetMetadata bool
 	breakList        bool
+	breakDestroy     bool
+	breakLease       bool
 
 	machines []*fly.Machine
+	leases   map[string]struct{}
 }
 
 func (m *mockFlapsClient) AcquireLease(ctx context.Context, machineID string, ttl *int) (*fly.MachineLease, error) {
-	return nil, fmt.Errorf("failed to acquire lease for %s", machineID)
+	if m.breakLease {
+		return nil, fmt.Errorf("failed to acquire lease for %s", machineID)
+	}
+
+	if m.leases == nil {
+		m.leases = make(map[string]struct{})
+	}
+	m.leases[machineID] = struct{}{}
+
+	return &fly.MachineLease{
+		Status: "success",
+		Data:   &fly.MachineLeaseData{Nonce: "nonce"},
+	}, nil
 }
 
 func (m *mockFlapsClient) Cordon(ctx context.Context, machineID string, nonce string) (err error) {
@@ -63,7 +82,10 @@ func (m *mockFlapsClient) DeleteVolume(ctx context.Context, volumeId string) (*f
 }
 
 func (m *mockFlapsClient) Destroy(ctx context.Context, input fly.RemoveMachineInput, nonce string) (err error) {
-	return fmt.Errorf("failed to destroy %s", input.ID)
+	if m.breakDestroy {
+		return fmt.Errorf("failed to destroy %s", input.ID)
+	}
+	return nil
 }
 
 func (m *mockFlapsClient) Exec(ctx context.Context, machineID string, in *fly.MachineExecRequest) (*fly.MachineExecResponse, error) {
@@ -122,13 +144,15 @@ func (m *mockFlapsClient) Launch(ctx context.Context, builder fly.LaunchMachineI
 	if m.breakLaunch {
 		return nil, fmt.Errorf("failed to launch %s", builder.ID)
 	}
-	return &fly.Machine{}, nil
+	globalMachineID += 1
+	return &fly.Machine{ID: fmt.Sprintf("%x", globalMachineID)}, nil
 }
 
 func (m *mockFlapsClient) List(ctx context.Context, state string) ([]*fly.Machine, error) {
 	if m.breakList {
 		return nil, fmt.Errorf("failed to list machines")
 	}
+	fmt.Printf("list\n")
 	return m.machines, nil
 }
 
@@ -149,11 +173,25 @@ func (m *mockFlapsClient) NewRequest(ctx context.Context, method, path string, i
 }
 
 func (m *mockFlapsClient) RefreshLease(ctx context.Context, machineID string, ttl *int, nonce string) (*fly.MachineLease, error) {
-	return nil, fmt.Errorf("failed to refresh lease for %s", machineID)
+	_, exists := m.leases[machineID]
+	if !exists {
+		return nil, fmt.Errorf("failed to refresh lease for %s", machineID)
+	}
+
+	return &fly.MachineLease{
+		Status: "success",
+		Data:   &fly.MachineLeaseData{Nonce: "nonce"},
+	}, nil
 }
 
 func (m *mockFlapsClient) ReleaseLease(ctx context.Context, machineID, nonce string) error {
-	return fmt.Errorf("failed to release lease for %s", machineID)
+	_, exists := m.leases[machineID]
+	if !exists {
+		return fmt.Errorf("failed to release lease for %s", machineID)
+	}
+	delete(m.leases, machineID)
+
+	return nil
 }
 
 func (m *mockFlapsClient) Restart(ctx context.Context, in fly.RestartMachineInput, nonce string) (err error) {

diff --git a/internal/command/deploy/plan.go b/internal/command/deploy/plan.go
@@ -369,6 +369,9 @@ func (md *machineDeployment) acquireLeases(ctx context.Context, machineTuples []
 	ctx, span := tracing.GetTracer().Start(ctx, "acquire_leases")
 
 	leaseGroup := errgroup.Group{}
+	if poolSize <= 0 {
+		panic("pool size must be > 0")
+	}
 	leaseGroup.SetLimit(poolSize)
 
 	for _, machineTuple := range machineTuples {

diff --git a/internal/command/deploy/web_client.go b/internal/command/deploy/web_client.go
@@ -0,0 +1,28 @@
+package deploy
+
+import (
+	"context"
+	"net"
+
+	fly "github.com/superfly/fly-go"
+	"github.com/superfly/flyctl/logs"
+)
+
+// webClient is a subset of web API that is needed for the deploy package.
+type webClient interface {
+	AddCertificate(ctx context.Context, appName, hostname string) (*fly.AppCertificate, *fly.HostnameCheck, error)
+	AllocateIPAddress(ctx context.Context, appName string, addrType string, region string, org *fly.Organization, network string) (*fly.IPAddress, error)
+	GetIPAddresses(ctx context.Context, appName string) ([]fly.IPAddress, error)
+	AllocateSharedIPAddress(ctx context.Context, appName string) (net.IP, error)
+
+	LatestImage(ctx context.Context, appName string) (string, error)
+
+	CreateRelease(ctx context.Context, input fly.CreateReleaseInput) (*fly.CreateReleaseResponse, error)
+	UpdateRelease(ctx context.Context, input fly.UpdateReleaseInput) (*fly.UpdateReleaseResponse, error)
+
+	GetApp(ctx context.Context, appName string) (*fly.App, error)
+	GetOrganizationBySlug(ctx context.Context, slug string) (*fly.Organization, error)
+
+	logs.WebClient
+	blueGreenWebClient
+}
diff --git a/internal/wireguard/wg.go b/internal/wireguard/wg.go
@@ -24,6 +24,10 @@ import (
 
 var cleanDNSPattern = regexp.MustCompile(`[^a-zA-Z0-9\\-]`)
 
+type WebClient interface {
+	ValidateWireGuardPeers(ctx context.Context, peerIPs []string) (invalid []string, err error)
+}
+
 func generatePeerName(ctx context.Context, apiClient flyutil.Client) (string, error) {
 	user, err := apiClient.GetCurrentUser(ctx)
 	if err != nil {
@@ -181,7 +185,7 @@ func setWireGuardStateForOrg(ctx context.Context, orgSlug, network string, s *wg
 	return setWireGuardState(ctx, states)
 }
 
-func PruneInvalidPeers(ctx context.Context, apiClient flyutil.Client) error {
+func PruneInvalidPeers(ctx context.Context, apiClient WebClient) error {
 	state, err := GetWireGuardState()
 	if err != nil {
 		return nil

diff --git a/logs/logs.go b/logs/logs.go
@@ -5,6 +5,9 @@ import (
 	"fmt"
 	"io"
 	"time"
+
+	fly "github.com/superfly/fly-go"
+	"github.com/superfly/flyctl/internal/wireguard"
 )
 
 type LogOptions struct {
@@ -17,6 +20,12 @@ type LogOptions struct {
 	NoTail     bool
 }
 
+type WebClient interface {
+	GetAppBasic(ctx context.Context, appName string) (*fly.AppBasic, error)
+	GetAppLogs(ctx context.Context, appName, token, region, instanceID string) (entries []fly.LogEntry, nextToken string, err error)
+	wireguard.WebClient
+}
+
 func (opts *LogOptions) toNatsSubject() (subject string) {
 	subject = fmt.Sprintf("logs.%s", opts.AppName)
 

diff --git a/logs/nats.go b/logs/nats.go
@@ -10,15 +10,14 @@ import (
 
 	"github.com/superfly/flyctl/agent"
 	"github.com/superfly/flyctl/internal/config"
-	"github.com/superfly/flyctl/internal/flyutil"
 )
 
 type natsLogStream struct {
 	nc  *nats.Conn
 	err error
 }
 
-func NewNatsStream(ctx context.Context, apiClient flyutil.Client, opts *LogOptions) (LogStream, error) {
+func NewNatsStream(ctx context.Context, apiClient WebClient, opts *LogOptions) (LogStream, error) {
 	app, err := apiClient.GetAppBasic(ctx, opts.AppName)
 	if err != nil {
 		return nil, fmt.Errorf("failed fetching target app: %w", err)

diff --git a/logs/polling.go b/logs/polling.go
@@ -8,15 +8,14 @@ import (
 	"github.com/pkg/errors"
 
 	fly "github.com/superfly/fly-go"
-	"github.com/superfly/flyctl/internal/flyutil"
 )
 
 type pollingStream struct {
 	err       error
-	apiClient flyutil.Client
+	apiClient WebClient
 }
 
-func NewPollingStream(client flyutil.Client) LogStream {
+func NewPollingStream(client WebClient) LogStream {
 	return &pollingStream{apiClient: client}
 }
 
@@ -36,7 +35,7 @@ func (s *pollingStream) Err() error {
 	return s.err
 }
 
-func Poll(ctx context.Context, out chan<- LogEntry, client flyutil.Client, opts *LogOptions) error {
+func Poll(ctx context.Context, out chan<- LogEntry, client WebClient, opts *LogOptions) error {
 	const (
 		minWait = time.Millisecond << 6
 		maxWait = minWait << 6