From 9f9e21dda693de6f3cec16f72a105afad4c416bf Mon Sep 17 00:00:00 2001 From: Zygimantas Date: Thu, 6 Feb 2025 20:51:55 -0500 Subject: [PATCH] fixe --- core/provider/digitalocean/droplet.go | 29 -------------------- core/provider/digitalocean/provider.go | 4 --- core/provider/digitalocean/provider_test.go | 30 ++++++++------------- core/provider/digitalocean/task_test.go | 14 +++++----- 4 files changed, 19 insertions(+), 58 deletions(-) diff --git a/core/provider/digitalocean/droplet.go b/core/provider/digitalocean/droplet.go index eae0bb8..c520644 100644 --- a/core/provider/digitalocean/droplet.go +++ b/core/provider/digitalocean/droplet.go @@ -3,7 +3,6 @@ package digitalocean import ( "context" "fmt" - "strings" "time" "github.com/pkg/errors" @@ -58,34 +57,6 @@ func (p *Provider) CreateDroplet(ctx context.Context, definition provider.TaskDe return droplet, nil } -func (t *Task) waitForSSHClient(ctx context.Context) error { - start := time.Now() - - err := util.WaitForCondition(ctx, time.Second*600, time.Millisecond*300, func() (bool, error) { - _, err := t.getDropletSSHClient(ctx) - - if err != nil { - if strings.Contains(err.Error(), "connection refused") { - t.logger.Debug("connection refused", zap.String("task", t.GetState().Name)) - return false, nil - } - return false, err - } - - return true, nil - }) - - if err != nil { - return errors.Wrap(err, "failed to wait for ssh in droplet to become active") - } - - end := time.Now() - - t.logger.Info("droplet's ssh daemon is ready after", zap.String("name", t.GetState().Name), zap.Duration("startup_time", end.Sub(start))) - - return nil -} - func (t *Task) waitForDockerStart(ctx context.Context) error { start := time.Now() diff --git a/core/provider/digitalocean/provider.go b/core/provider/digitalocean/provider.go index 26e1190..fc765de 100644 --- a/core/provider/digitalocean/provider.go +++ b/core/provider/digitalocean/provider.go @@ -163,10 +163,6 @@ func (p *Provider) CreateTask(ctx context.Context, definition provider.TaskDefin return nil, fmt.Errorf("failed to wait for tailscale peer: %w", err) } - if err := task.waitForSSHClient(ctx); err != nil { - return nil, fmt.Errorf("failed to wait for SSH client: %w", err) - } - ip, err := task.GetIP(ctx) if err != nil { diff --git a/core/provider/digitalocean/provider_test.go b/core/provider/digitalocean/provider_test.go index 4275a6c..a96e46c 100644 --- a/core/provider/digitalocean/provider_test.go +++ b/core/provider/digitalocean/provider_test.go @@ -52,8 +52,6 @@ func setupTestProvider(t *testing.T, ctx context.Context) (*Provider, *mocks.Moc mockTailscaleServer := clientmocks.NewMockTailscaleServer(t) mockTailscaleClient := clientmocks.NewMockTailscaleLocalClient(t) - mockTailscaleClient.On("Status", ctx).Return(generateTailscaleStatus(t, "test-task", "1.2.3.4"), nil) - mockTailscale := TailscaleSettings{ Server: mockTailscaleServer, LocalClient: mockTailscaleClient, @@ -87,8 +85,6 @@ func setupTestProvider(t *testing.T, ctx context.Context) (*Provider, *mocks.Moc mockDO.On("CreateTag", ctx, mock.Anything).Return(&godo.Tag{Name: "test-tag"}, nil) mockDO.On("CreateFirewall", ctx, mock.Anything).Return(&godo.Firewall{ID: "test-firewall"}, nil) - mockDO.On("GetKeyByFingerprint", ctx, mock.AnythingOfType("string")).Return(nil, nil) - mockDO.On("CreateKey", ctx, mock.Anything).Return(&godo.Key{}, nil) mockDockerClients := map[string]clients.DockerClient{ "test-task": mockDocker, @@ -97,6 +93,8 @@ func setupTestProvider(t *testing.T, ctx context.Context) (*Provider, *mocks.Moc p, err := NewProviderWithClient(ctx, "test-provider", mockDO, mockTailscale, WithDockerClients(mockDockerClients), WithLogger(logger)) require.NoError(t, err) + mockTailscaleClient.On("Status", ctx).Return(generateTailscaleStatus(t, fmt.Sprintf("%s-test-task", p.GetState().PetriTag), "1.2.3.4"), nil) + droplet := &godo.Droplet{ ID: 123, Networks: &godo.Networks{ @@ -181,8 +179,6 @@ func setupValidationTestProvider(t *testing.T, ctx context.Context) *Provider { mockDO.On("CreateTag", ctx, mock.Anything).Return(&godo.Tag{Name: "test-tag"}, nil) mockDO.On("CreateFirewall", ctx, mock.Anything).Return(&godo.Firewall{ID: "test-firewall"}, nil) - mockDO.On("GetKeyByFingerprint", ctx, mock.AnythingOfType("string")).Return(nil, nil) - mockDO.On("CreateKey", ctx, mock.Anything).Return(&godo.Key{}, nil) p, err := NewProviderWithClient(ctx, "test-provider", mockDO, mockTailscale, WithLogger(logger)) require.NoError(t, err) @@ -332,23 +328,22 @@ func TestConcurrentTaskCreationAndCleanup(t *testing.T) { }, nil).Maybe() mockDocker.On("Close").Return(nil).Once() - mockStatus := ipnstate.PeerStatus{HostName: fmt.Sprintf("test-task-%d", i), TailscaleIPs: []netip.Addr{netip.MustParseAddr(fmt.Sprintf("1.2.3.%d", i+1))}} - mockStatuses[key.NewNode().Public()] = &mockStatus } - mockTailscaleClient.On("Status", ctx).Return(&ipnstate.Status{Peer: mockStatuses}, nil) - mockDO.On("CreateTag", ctx, mock.Anything).Return(&godo.Tag{Name: "test-tag"}, nil) mockDO.On("CreateFirewall", ctx, mock.Anything).Return(&godo.Firewall{ID: "test-firewall"}, nil) - mockDO.On("GetKeyByFingerprint", ctx, mock.AnythingOfType("string")). - Return(nil, nil) - mockDO.On("CreateKey", ctx, mock.Anything).Return(&godo.Key{}, nil) - p, err := NewProviderWithClient(ctx, "test-provider", mockDO, mockTailscale, WithDockerClients(mockDockerClients), WithLogger(logger)) require.NoError(t, err) + for i := 0; i < 10; i++ { + mockStatus := ipnstate.PeerStatus{HostName: fmt.Sprintf("%s-test-task-%d", p.GetState().PetriTag, i), TailscaleIPs: []netip.Addr{netip.MustParseAddr(fmt.Sprintf("1.2.3.%d", i+1))}} + mockStatuses[key.NewNode().Public()] = &mockStatus + } + + mockTailscaleClient.On("Status", ctx).Return(&ipnstate.Status{Peer: mockStatuses}, nil) + numTasks := 10 var wg sync.WaitGroup errors := make(chan error, numTasks) @@ -380,7 +375,6 @@ func TestConcurrentTaskCreationAndCleanup(t *testing.T) { mockDO.On("DeleteDropletByTag", ctx, mock.AnythingOfType("string")).Return(nil).Once() mockDO.On("DeleteFirewall", ctx, mock.AnythingOfType("string")).Return(nil).Once() - mockDO.On("DeleteKeyByFingerprint", ctx, mock.AnythingOfType("string")).Return(nil).Once() mockDO.On("DeleteTag", ctx, mock.AnythingOfType("string")).Return(nil).Once() for i := 0; i < numTasks; i++ { @@ -504,8 +498,6 @@ func TestProviderSerialization(t *testing.T) { mockTailscaleServer := clientmocks.NewMockTailscaleServer(t) mockTailscaleClient := clientmocks.NewMockTailscaleLocalClient(t) - mockTailscaleClient.On("Status", ctx).Return(generateTailscaleStatus(t, "test-task", "1.2.3.4"), nil) - mockTailscale := TailscaleSettings{ Server: mockTailscaleServer, LocalClient: mockTailscaleClient, @@ -515,8 +507,6 @@ func TestProviderSerialization(t *testing.T) { mockDO.On("CreateTag", ctx, mock.Anything).Return(&godo.Tag{Name: "petri-droplet-test"}, nil) mockDO.On("CreateFirewall", ctx, mock.Anything).Return(&godo.Firewall{ID: "test-firewall"}, nil) - mockDO.On("GetKeyByFingerprint", ctx, mock.AnythingOfType("string")).Return(nil, nil) - mockDO.On("CreateKey", ctx, mock.Anything).Return(&godo.Key{}, nil) mockDockerClients := map[string]clients.DockerClient{ "test-task": mockDocker, @@ -525,6 +515,8 @@ func TestProviderSerialization(t *testing.T) { p1, err := NewProviderWithClient(ctx, "test-provider", mockDO, mockTailscale, WithDockerClients(mockDockerClients), WithLogger(zap.NewExample())) require.NoError(t, err) + mockTailscaleClient.On("Status", ctx).Return(generateTailscaleStatus(t, fmt.Sprintf("%s-test-task", p1.GetState().PetriTag), "1.2.3.4"), nil) + droplet := &godo.Droplet{ ID: 123, Networks: &godo.Networks{ diff --git a/core/provider/digitalocean/task_test.go b/core/provider/digitalocean/task_test.go index 351902f..6a17633 100644 --- a/core/provider/digitalocean/task_test.go +++ b/core/provider/digitalocean/task_test.go @@ -395,9 +395,10 @@ func TestTaskGetIP(t *testing.T) { task := &Task{ state: &TaskState{ - ID: strconv.Itoa(droplet.ID), - Name: "test-task", - ProviderName: "test-provider", + TailscaleHostname: "test-task", + ID: strconv.Itoa(droplet.ID), + Name: "test-task", + ProviderName: "test-provider", }, logger: logger, dockerClient: mockDocker, @@ -759,9 +760,10 @@ func TestTaskExposingPort(t *testing.T) { task := &Task{ state: &TaskState{ - ID: strconv.Itoa(droplet.ID), - Name: "test-task", - ProviderName: "test-provider", + TailscaleHostname: "test-task", + ID: strconv.Itoa(droplet.ID), + Name: "test-task", + ProviderName: "test-provider", Definition: provider.TaskDefinition{ Name: "test-task", Image: provider.ImageDefinition{