From da8f06d30629a409a6dc7bdbd24ca57c2f21c69e Mon Sep 17 00:00:00 2001 From: David Koenitzer Date: Tue, 17 Oct 2023 14:50:17 -0400 Subject: [PATCH 01/13] fix runtime validation --- astro-client/astro.go | 14 +++++++++++ astro-client/mocks/Client.go | 24 ++++++++++++++++++ astro-client/queries.go | 45 ++++++++++++++++++++++++++++++++++ cloud/deployment/deployment.go | 11 ++++++--- 4 files changed, 91 insertions(+), 3 deletions(-) diff --git a/astro-client/astro.go b/astro-client/astro.go index 2461a5e46..74a8a718d 100644 --- a/astro-client/astro.go +++ b/astro-client/astro.go @@ -21,6 +21,7 @@ type Client interface { DeleteDeployment(input DeleteDeploymentInput) (Deployment, error) GetDeploymentHistory(vars map[string]interface{}) (DeploymentHistory, error) GetDeploymentConfig() (DeploymentConfig, error) + GetDeploymentConfigWithOrganization(organizationID string) (DeploymentConfig, error) ModifyDeploymentVariable(input EnvironmentVariablesInput) ([]EnvironmentVariablesObject, error) InitiateDagDeployment(input InitiateDagDeploymentInput) (InitiateDagDeployment, error) ReportDagDeploymentStatus(input *ReportDagDeploymentStatusInput) (DagDeploymentStatus, error) @@ -125,6 +126,19 @@ func (c *HTTPClient) GetDeploymentConfig() (DeploymentConfig, error) { return resp.Data.GetDeploymentConfig, nil } +func (c *HTTPClient) GetDeploymentConfigWithOrganization(organizationID string) (DeploymentConfig, error) { + req := Request{ + Query: GetDeploymentConfigOptionsWithOrganization, + Variables: map[string]interface{}{"organizationId": organizationID}, + } + + resp, err := req.DoWithPublicClient(c) + if err != nil { + return DeploymentConfig{}, err + } + return resp.Data.GetDeploymentConfig, nil +} + func (c *HTTPClient) ModifyDeploymentVariable(input EnvironmentVariablesInput) ([]EnvironmentVariablesObject, error) { req := Request{ Query: CreateDeploymentVariables, diff --git a/astro-client/mocks/Client.go b/astro-client/mocks/Client.go index 90680be96..5025eb7ae 100644 --- a/astro-client/mocks/Client.go +++ b/astro-client/mocks/Client.go @@ -163,6 +163,30 @@ func (_m *Client) GetDeploymentConfig() (astro.DeploymentConfig, error) { return r0, r1 } +// GetDeploymentConfigWithOrganization provides a mock function with given fields: organizationID +func (_m *Client) GetDeploymentConfigWithOrganization(organizationID string) (astro.DeploymentConfig, error) { + ret := _m.Called(organizationID) + + var r0 astro.DeploymentConfig + var r1 error + if rf, ok := ret.Get(0).(func(string) (astro.DeploymentConfig, error)); ok { + return rf(organizationID) + } + if rf, ok := ret.Get(0).(func(string) astro.DeploymentConfig); ok { + r0 = rf(organizationID) + } else { + r0 = ret.Get(0).(astro.DeploymentConfig) + } + + if rf, ok := ret.Get(1).(func(string) error); ok { + r1 = rf(organizationID) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + // GetDeploymentHistory provides a mock function with given fields: vars func (_m *Client) GetDeploymentHistory(vars map[string]interface{}) (astro.DeploymentHistory, error) { ret := _m.Called(vars) diff --git a/astro-client/queries.go b/astro-client/queries.go index 93dc538af..b0d4bda0b 100644 --- a/astro-client/queries.go +++ b/astro-client/queries.go @@ -208,6 +208,51 @@ var ( } ` + GetDeploymentConfigOptionsWithOrganization = ` + query deploymentConfigOptions($organizationId: Id!) { + deploymentConfigOptions(organizationId: $organizationId) { + components + astroUnit { + cpu + memory + } + executors + runtimeReleases { + channel + version + } + astroMachines { + concurrentTasks + concurrentTasksMax + cpu + memory + nodePoolType + storageSize + type + } + defaultAstroMachine { + concurrentTasks + concurrentTasksMax + cpu + memory + nodePoolType + storageSize + type + } + defaultSchedulerSize { + cpu + memory + size + } + schedulerSizes { + cpu + memory + size + } + } + } + ` + GetWorkerQueueOptions = ` query workerQueueOptions { workerQueueOptions { diff --git a/cloud/deployment/deployment.go b/cloud/deployment/deployment.go index c9a9c684e..31c2760d8 100644 --- a/cloud/deployment/deployment.go +++ b/cloud/deployment/deployment.go @@ -383,7 +383,12 @@ func validateResources(schedulerAU, schedulerReplicas int, configOption astro.De } func validateRuntimeVersion(runtimeVersion string, client astro.Client) (bool, error) { - runtimeReleases, err := GetRuntimeReleases(client) + c, err := config.GetCurrentContext() + if err != nil { + return false, err + } + + runtimeReleases, err := GetRuntimeReleases(c.Organization, client) if err != nil { return false, err } @@ -394,11 +399,11 @@ func validateRuntimeVersion(runtimeVersion string, client astro.Client) (bool, e return true, nil } -func GetRuntimeReleases(client astro.Client) ([]string, error) { +func GetRuntimeReleases(organizationID string, client astro.Client) ([]string, error) { // get deployment config options runtimeReleases := []string{} - ConfigOptions, err := client.GetDeploymentConfig() + ConfigOptions, err := client.GetDeploymentConfigWithOrganization(organizationID) if err != nil { return runtimeReleases, errors.Wrap(err, astro.AstronomerConnectionErrMsg) } From 023f973e069cd90ca47e69497200bb486423ca1c Mon Sep 17 00:00:00 2001 From: David Koenitzer Date: Tue, 17 Oct 2023 14:55:58 -0400 Subject: [PATCH 02/13] fix runtime validation --- cmd/cloud/deployment_test.go | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/cmd/cloud/deployment_test.go b/cmd/cloud/deployment_test.go index 269285bfa..f17bee873 100644 --- a/cmd/cloud/deployment_test.go +++ b/cmd/cloud/deployment_test.go @@ -173,6 +173,26 @@ func TestDeploymentCreate(t *testing.T) { }, }, }, nil).Times(14) + mockClient.On("GetDeploymentConfigWithOrganization", mock.Anything).Return(astro.DeploymentConfig{ + Components: astro.Components{ + Scheduler: astro.SchedulerConfig{ + AU: astro.AuConfig{ + Default: 5, + Limit: 24, + }, + Replicas: astro.ReplicasConfig{ + Default: 1, + Minimum: 1, + Limit: 4, + }, + }, + }, + RuntimeReleases: []astro.RuntimeRelease{ + { + Version: "4.2.5", + }, + }, + }, nil).Times(14) mockCoreClient.On("ListWorkspacesWithResponse", mock.Anything, mock.Anything, mock.Anything).Return(&ListWorkspacesResponseOK, nil).Times(5) mockCoreClient.On("ListClustersWithResponse", mock.Anything, mock.Anything, mock.Anything).Return(&mockListClustersResponse, nil).Times(4) mockClient.On("CreateDeployment", &deploymentCreateInput).Return(astro.Deployment{ID: "test-id"}, nil).Twice() From 21ec9b6a357b3ac3a4e1d0fafa439a573f29ab7b Mon Sep 17 00:00:00 2001 From: David Koenitzer Date: Tue, 17 Oct 2023 15:12:46 -0400 Subject: [PATCH 03/13] fix test --- cmd/cloud/deployment_test.go | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/cmd/cloud/deployment_test.go b/cmd/cloud/deployment_test.go index f17bee873..1d88fc234 100644 --- a/cmd/cloud/deployment_test.go +++ b/cmd/cloud/deployment_test.go @@ -172,7 +172,7 @@ func TestDeploymentCreate(t *testing.T) { Version: "4.2.5", }, }, - }, nil).Times(14) + }, nil).Times(7) mockClient.On("GetDeploymentConfigWithOrganization", mock.Anything).Return(astro.DeploymentConfig{ Components: astro.Components{ Scheduler: astro.SchedulerConfig{ @@ -192,7 +192,7 @@ func TestDeploymentCreate(t *testing.T) { Version: "4.2.5", }, }, - }, nil).Times(14) + }, nil).Times(7) mockCoreClient.On("ListWorkspacesWithResponse", mock.Anything, mock.Anything, mock.Anything).Return(&ListWorkspacesResponseOK, nil).Times(5) mockCoreClient.On("ListClustersWithResponse", mock.Anything, mock.Anything, mock.Anything).Return(&mockListClustersResponse, nil).Times(4) mockClient.On("CreateDeployment", &deploymentCreateInput).Return(astro.Deployment{ID: "test-id"}, nil).Twice() From 61dfe941647c6eb6020660819233d259440d9c9a Mon Sep 17 00:00:00 2001 From: David Koenitzer Date: Tue, 17 Oct 2023 15:21:40 -0400 Subject: [PATCH 04/13] fix test --- cloud/deployment/deployment_test.go | 22 +++++++++++++++++++++- 1 file changed, 21 insertions(+), 1 deletion(-) diff --git a/cloud/deployment/deployment_test.go b/cloud/deployment/deployment_test.go index 06fef339b..5667dd901 100644 --- a/cloud/deployment/deployment_test.go +++ b/cloud/deployment/deployment_test.go @@ -779,7 +779,27 @@ func TestCreate(t *testing.T) { Version: "4.2.5", }, }, - }, nil).Times(2) + }, nil).Times(1) + mockClient.On("GetDeploymentConfigWithOrganization", mock.Anything).Return(astro.DeploymentConfig{ + Components: astro.Components{ + Scheduler: astro.SchedulerConfig{ + AU: astro.AuConfig{ + Default: 5, + Limit: 24, + }, + Replicas: astro.ReplicasConfig{ + Default: 1, + Minimum: 1, + Limit: 4, + }, + }, + }, + RuntimeReleases: []astro.RuntimeRelease{ + { + Version: "4.2.5", + }, + }, + }, nil).Times(1) mockCoreClient.On("ListWorkspacesWithResponse", mock.Anything, mock.Anything, mock.Anything).Return(&ListWorkspacesResponseOK, nil).Once() mockCoreClient.On("ListClustersWithResponse", mock.Anything, mockOrgShortName, clusterListParams).Return(&mockListClustersResponse, nil).Once() mockClient.On("CreateDeployment", &deploymentCreateInput).Return(astro.Deployment{ID: "test-id"}, nil).Once() From 58389b804e6773cce45fed88d24ee82ac80cf925 Mon Sep 17 00:00:00 2001 From: David Koenitzer Date: Tue, 17 Oct 2023 15:26:40 -0400 Subject: [PATCH 05/13] fix test --- cloud/deployment/deployment_test.go | 172 ++++++++++++++++++++++++++-- 1 file changed, 165 insertions(+), 7 deletions(-) diff --git a/cloud/deployment/deployment_test.go b/cloud/deployment/deployment_test.go index 5667dd901..072e53f7f 100644 --- a/cloud/deployment/deployment_test.go +++ b/cloud/deployment/deployment_test.go @@ -831,7 +831,27 @@ func TestCreate(t *testing.T) { Version: "4.2.5", }, }, - }, nil).Times(2) + }, nil).Times(1) + mockClient.On("GetDeploymentConfigWithOrganization", mock.Anything).Return(astro.DeploymentConfig{ + Components: astro.Components{ + Scheduler: astro.SchedulerConfig{ + AU: astro.AuConfig{ + Default: 5, + Limit: 24, + }, + Replicas: astro.ReplicasConfig{ + Default: 1, + Minimum: 1, + Limit: 4, + }, + }, + }, + RuntimeReleases: []astro.RuntimeRelease{ + { + Version: "4.2.5", + }, + }, + }, nil).Times(1) deploymentCreateInput.APIKeyOnlyDeployments = true defer func() { deploymentCreateInput.APIKeyOnlyDeployments = false }() mockCoreClient.On("ListWorkspacesWithResponse", mock.Anything, mock.Anything, mock.Anything).Return(&ListWorkspacesResponseOK, nil).Once() @@ -876,7 +896,30 @@ func TestCreate(t *testing.T) { DefaultSchedulerSize: astro.MachineUnit{ Size: "small", }, - }, nil).Times(2) + }, nil).Times(1) + mockClient.On("GetDeploymentConfigWithOrganization", mock.Anything).Return(astro.DeploymentConfig{ + Components: astro.Components{ + Scheduler: astro.SchedulerConfig{ + AU: astro.AuConfig{ + Default: 5, + Limit: 24, + }, + Replicas: astro.ReplicasConfig{ + Default: 1, + Minimum: 1, + Limit: 4, + }, + }, + }, + RuntimeReleases: []astro.RuntimeRelease{ + { + Version: "4.2.5", + }, + }, + DefaultSchedulerSize: astro.MachineUnit{ + Size: "small", + }, + }, nil).Times(1) getSharedClusterParams := &astrocore.GetSharedClusterParams{ Region: region, CloudProvider: astrocore.GetSharedClusterParamsCloudProvider(astrocore.SharedClusterCloudProviderGcp), @@ -965,7 +1008,30 @@ func TestCreate(t *testing.T) { DefaultSchedulerSize: astro.MachineUnit{ Size: "small", }, - }, nil).Times(2) + }, nil).Times(1) + mockClient.On("GetDeploymentConfigWithOrgnaiztion", mock.Anything).Return(astro.DeploymentConfig{ + Components: astro.Components{ + Scheduler: astro.SchedulerConfig{ + AU: astro.AuConfig{ + Default: 5, + Limit: 24, + }, + Replicas: astro.ReplicasConfig{ + Default: 1, + Minimum: 1, + Limit: 4, + }, + }, + }, + RuntimeReleases: []astro.RuntimeRelease{ + { + Version: "4.2.5", + }, + }, + DefaultSchedulerSize: astro.MachineUnit{ + Size: "small", + }, + }, nil).Times(1) getSharedClusterParams := &astrocore.GetSharedClusterParams{ Region: region, CloudProvider: astrocore.GetSharedClusterParamsCloudProvider(astrocore.SharedClusterCloudProviderGcp), @@ -1046,7 +1112,30 @@ func TestCreate(t *testing.T) { Version: "4.2.5", }, }, - }, nil).Times(2) + }, nil).Times(1) + mockClient.On("GetDeploymentConfigWithOrgnaiztion", mock.Anything).Return(astro.DeploymentConfig{ + Components: astro.Components{ + Scheduler: astro.SchedulerConfig{ + AU: astro.AuConfig{ + Default: 5, + Limit: 24, + }, + Replicas: astro.ReplicasConfig{ + Default: 1, + Minimum: 1, + Limit: 4, + }, + }, + }, + RuntimeReleases: []astro.RuntimeRelease{ + { + Version: "4.2.5", + }, + }, + DefaultSchedulerSize: astro.MachineUnit{ + Size: "small", + }, + }, nil).Times(1) deploymentCreateInput.DeploymentSpec.Executor = "KubeExecutor" defer func() { deploymentCreateInput.DeploymentSpec.Executor = CeleryExecutor }() mockCoreClient.On("ListWorkspacesWithResponse", mock.Anything, mock.Anything, mock.Anything).Return(&ListWorkspacesResponseOK, nil).Once() @@ -1081,7 +1170,30 @@ func TestCreate(t *testing.T) { Version: "4.2.5", }, }, - }, nil).Times(4) + }, nil).Times(2) + mockClient.On("GetDeploymentConfigWithOrgnaiztion", mock.Anything).Return(astro.DeploymentConfig{ + Components: astro.Components{ + Scheduler: astro.SchedulerConfig{ + AU: astro.AuConfig{ + Default: 5, + Limit: 24, + }, + Replicas: astro.ReplicasConfig{ + Default: 1, + Minimum: 1, + Limit: 4, + }, + }, + }, + RuntimeReleases: []astro.RuntimeRelease{ + { + Version: "4.2.5", + }, + }, + DefaultSchedulerSize: astro.MachineUnit{ + Size: "small", + }, + }, nil).Times(2) mockCoreClient.On("ListWorkspacesWithResponse", mock.Anything, mock.Anything, mock.Anything).Return(&ListWorkspacesResponseOK, nil).Twice() mockCoreClient.On("ListClustersWithResponse", mock.Anything, mock.Anything, clusterListParams).Return(&mockListClustersResponse, nil).Twice() mockClient.On("CreateDeployment", &deploymentCreateInput).Return(astro.Deployment{ID: deploymentID}, nil).Twice() @@ -1141,7 +1253,30 @@ func TestCreate(t *testing.T) { Version: "4.2.5", }, }, - }, nil).Times(2) + }, nil).Times(1) + mockClient.On("GetDeploymentConfigWithOrgnaiztion", mock.Anything).Return(astro.DeploymentConfig{ + Components: astro.Components{ + Scheduler: astro.SchedulerConfig{ + AU: astro.AuConfig{ + Default: 5, + Limit: 24, + }, + Replicas: astro.ReplicasConfig{ + Default: 1, + Minimum: 1, + Limit: 4, + }, + }, + }, + RuntimeReleases: []astro.RuntimeRelease{ + { + Version: "4.2.5", + }, + }, + DefaultSchedulerSize: astro.MachineUnit{ + Size: "small", + }, + }, nil).Times(1) mockCoreClient.On("ListWorkspacesWithResponse", mock.Anything, mock.Anything, mock.Anything).Return(&ListWorkspacesResponseOK, nil).Once() mockCoreClient.On("ListClustersWithResponse", mock.Anything, mock.Anything, clusterListParams).Return(&mockListClustersResponse, nil).Once() mockClient.On("CreateDeployment", &deploymentCreateInput).Return(astro.Deployment{}, errMock).Once() @@ -1181,7 +1316,30 @@ func TestCreate(t *testing.T) { Version: "4.2.5", }, }, - }, nil).Times(2) + }, nil).Times(1) + mockClient.On("GetDeploymentConfigWithOrgnaiztion", mock.Anything).Return(astro.DeploymentConfig{ + Components: astro.Components{ + Scheduler: astro.SchedulerConfig{ + AU: astro.AuConfig{ + Default: 5, + Limit: 24, + }, + Replicas: astro.ReplicasConfig{ + Default: 1, + Minimum: 1, + Limit: 4, + }, + }, + }, + RuntimeReleases: []astro.RuntimeRelease{ + { + Version: "4.2.5", + }, + }, + DefaultSchedulerSize: astro.MachineUnit{ + Size: "small", + }, + }, nil).Times(1) mockCoreClient.On("ListWorkspacesWithResponse", mock.Anything, mock.Anything, mock.Anything).Return(&ListWorkspacesResponseOK, nil).Once() mockCoreClient.On("ListClustersWithResponse", mock.Anything, mock.Anything, clusterListParams).Return(&astrocore.ListClustersResponse{}, errMock).Once() err := Create("test-name", ws, "test-desc", "invalid-cluster-id", "4.2.5", dagDeploy, CeleryExecutor, "", "", "", "", "", 10, 3, mockClient, mockCoreClient, false, &disableCiCdEnforcement) From 9f5098491f22b09a80229f7ceb0ccb1d8aa02c71 Mon Sep 17 00:00:00 2001 From: David Koenitzer Date: Tue, 17 Oct 2023 15:46:24 -0400 Subject: [PATCH 06/13] fix test --- cloud/deployment/deployment_test.go | 25 +------------------------ 1 file changed, 1 insertion(+), 24 deletions(-) diff --git a/cloud/deployment/deployment_test.go b/cloud/deployment/deployment_test.go index 072e53f7f..47576ee32 100644 --- a/cloud/deployment/deployment_test.go +++ b/cloud/deployment/deployment_test.go @@ -986,29 +986,6 @@ func TestCreate(t *testing.T) { ctx.SetContextKey("organization_product", "HOSTED") ctx.SetContextKey("organization", org) ctx.SetContextKey("workspace", ws) - mockClient.On("GetDeploymentConfig").Return(astro.DeploymentConfig{ - Components: astro.Components{ - Scheduler: astro.SchedulerConfig{ - AU: astro.AuConfig{ - Default: 5, - Limit: 24, - }, - Replicas: astro.ReplicasConfig{ - Default: 1, - Minimum: 1, - Limit: 4, - }, - }, - }, - RuntimeReleases: []astro.RuntimeRelease{ - { - Version: "4.2.5", - }, - }, - DefaultSchedulerSize: astro.MachineUnit{ - Size: "small", - }, - }, nil).Times(1) mockClient.On("GetDeploymentConfigWithOrgnaiztion", mock.Anything).Return(astro.DeploymentConfig{ Components: astro.Components{ Scheduler: astro.SchedulerConfig{ @@ -1031,7 +1008,7 @@ func TestCreate(t *testing.T) { DefaultSchedulerSize: astro.MachineUnit{ Size: "small", }, - }, nil).Times(1) + }, nil).Times(2) getSharedClusterParams := &astrocore.GetSharedClusterParams{ Region: region, CloudProvider: astrocore.GetSharedClusterParamsCloudProvider(astrocore.SharedClusterCloudProviderGcp), From c5857b02809496c23526b0616513773f9d728216 Mon Sep 17 00:00:00 2001 From: David Koenitzer Date: Tue, 17 Oct 2023 17:13:12 -0400 Subject: [PATCH 07/13] fix test --- cloud/deployment/deployment_test.go | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/cloud/deployment/deployment_test.go b/cloud/deployment/deployment_test.go index 47576ee32..82ceca67f 100644 --- a/cloud/deployment/deployment_test.go +++ b/cloud/deployment/deployment_test.go @@ -986,6 +986,29 @@ func TestCreate(t *testing.T) { ctx.SetContextKey("organization_product", "HOSTED") ctx.SetContextKey("organization", org) ctx.SetContextKey("workspace", ws) + mockClient.On("GetDeploymentConfig").Return(astro.DeploymentConfig{ + Components: astro.Components{ + Scheduler: astro.SchedulerConfig{ + AU: astro.AuConfig{ + Default: 5, + Limit: 24, + }, + Replicas: astro.ReplicasConfig{ + Default: 1, + Minimum: 1, + Limit: 4, + }, + }, + }, + RuntimeReleases: []astro.RuntimeRelease{ + { + Version: "4.2.5", + }, + }, + DefaultSchedulerSize: astro.MachineUnit{ + Size: "small", + }, + }, nil).Times(2) mockClient.On("GetDeploymentConfigWithOrgnaiztion", mock.Anything).Return(astro.DeploymentConfig{ Components: astro.Components{ Scheduler: astro.SchedulerConfig{ From 3734c11ab3f1e7e54c65395bf184a35fc5f540e6 Mon Sep 17 00:00:00 2001 From: David Koenitzer Date: Tue, 17 Oct 2023 17:18:35 -0400 Subject: [PATCH 08/13] fix test --- cloud/deployment/deployment_test.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cloud/deployment/deployment_test.go b/cloud/deployment/deployment_test.go index 82ceca67f..1eca58dce 100644 --- a/cloud/deployment/deployment_test.go +++ b/cloud/deployment/deployment_test.go @@ -1031,7 +1031,7 @@ func TestCreate(t *testing.T) { DefaultSchedulerSize: astro.MachineUnit{ Size: "small", }, - }, nil).Times(2) + }, nil).Times(3) getSharedClusterParams := &astrocore.GetSharedClusterParams{ Region: region, CloudProvider: astrocore.GetSharedClusterParamsCloudProvider(astrocore.SharedClusterCloudProviderGcp), From 4243ebf5d0eb07f8d1a3d3f885db61ffa940992d Mon Sep 17 00:00:00 2001 From: David Koenitzer Date: Wed, 18 Oct 2023 15:58:44 -0400 Subject: [PATCH 09/13] fix spelling error --- cloud/deployment/deployment_test.go | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/cloud/deployment/deployment_test.go b/cloud/deployment/deployment_test.go index 1eca58dce..9bb7cc55a 100644 --- a/cloud/deployment/deployment_test.go +++ b/cloud/deployment/deployment_test.go @@ -1009,7 +1009,7 @@ func TestCreate(t *testing.T) { Size: "small", }, }, nil).Times(2) - mockClient.On("GetDeploymentConfigWithOrgnaiztion", mock.Anything).Return(astro.DeploymentConfig{ + mockClient.On("GetDeploymentConfigWithOrganization", mock.Anything).Return(astro.DeploymentConfig{ Components: astro.Components{ Scheduler: astro.SchedulerConfig{ AU: astro.AuConfig{ @@ -1031,7 +1031,7 @@ func TestCreate(t *testing.T) { DefaultSchedulerSize: astro.MachineUnit{ Size: "small", }, - }, nil).Times(3) + }, nil).Times(2) getSharedClusterParams := &astrocore.GetSharedClusterParams{ Region: region, CloudProvider: astrocore.GetSharedClusterParamsCloudProvider(astrocore.SharedClusterCloudProviderGcp), @@ -1113,7 +1113,7 @@ func TestCreate(t *testing.T) { }, }, }, nil).Times(1) - mockClient.On("GetDeploymentConfigWithOrgnaiztion", mock.Anything).Return(astro.DeploymentConfig{ + mockClient.On("GetDeploymentConfigWithOrganization", mock.Anything).Return(astro.DeploymentConfig{ Components: astro.Components{ Scheduler: astro.SchedulerConfig{ AU: astro.AuConfig{ @@ -1171,7 +1171,7 @@ func TestCreate(t *testing.T) { }, }, }, nil).Times(2) - mockClient.On("GetDeploymentConfigWithOrgnaiztion", mock.Anything).Return(astro.DeploymentConfig{ + mockClient.On("GetDeploymentConfigWithOrganization", mock.Anything).Return(astro.DeploymentConfig{ Components: astro.Components{ Scheduler: astro.SchedulerConfig{ AU: astro.AuConfig{ @@ -1254,7 +1254,7 @@ func TestCreate(t *testing.T) { }, }, }, nil).Times(1) - mockClient.On("GetDeploymentConfigWithOrgnaiztion", mock.Anything).Return(astro.DeploymentConfig{ + mockClient.On("GetDeploymentConfigWithOrganization", mock.Anything).Return(astro.DeploymentConfig{ Components: astro.Components{ Scheduler: astro.SchedulerConfig{ AU: astro.AuConfig{ @@ -1317,7 +1317,7 @@ func TestCreate(t *testing.T) { }, }, }, nil).Times(1) - mockClient.On("GetDeploymentConfigWithOrgnaiztion", mock.Anything).Return(astro.DeploymentConfig{ + mockClient.On("GetDeploymentConfigWithOrganization", mock.Anything).Return(astro.DeploymentConfig{ Components: astro.Components{ Scheduler: astro.SchedulerConfig{ AU: astro.AuConfig{ From 998e0d1f09b2d60da2c4f03bf824c9ddc0bfbfce Mon Sep 17 00:00:00 2001 From: David Koenitzer Date: Thu, 19 Oct 2023 10:18:35 -0400 Subject: [PATCH 10/13] fix test --- cloud/deployment/deployment_test.go | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/cloud/deployment/deployment_test.go b/cloud/deployment/deployment_test.go index 9bb7cc55a..4accd38b6 100644 --- a/cloud/deployment/deployment_test.go +++ b/cloud/deployment/deployment_test.go @@ -1008,7 +1008,7 @@ func TestCreate(t *testing.T) { DefaultSchedulerSize: astro.MachineUnit{ Size: "small", }, - }, nil).Times(2) + }, nil).Times(1) mockClient.On("GetDeploymentConfigWithOrganization", mock.Anything).Return(astro.DeploymentConfig{ Components: astro.Components{ Scheduler: astro.SchedulerConfig{ @@ -1031,7 +1031,7 @@ func TestCreate(t *testing.T) { DefaultSchedulerSize: astro.MachineUnit{ Size: "small", }, - }, nil).Times(2) + }, nil).Times(1) getSharedClusterParams := &astrocore.GetSharedClusterParams{ Region: region, CloudProvider: astrocore.GetSharedClusterParamsCloudProvider(astrocore.SharedClusterCloudProviderGcp), From 15ce10d50745135b09beda7fd24cdec64328bdcd Mon Sep 17 00:00:00 2001 From: David Koenitzer Date: Thu, 19 Oct 2023 10:35:24 -0400 Subject: [PATCH 11/13] fix test --- cloud/deployment/deployment_test.go | 88 +++++++++++++++++++++++++++-- 1 file changed, 84 insertions(+), 4 deletions(-) diff --git a/cloud/deployment/deployment_test.go b/cloud/deployment/deployment_test.go index 429babfe0..96da60fba 100644 --- a/cloud/deployment/deployment_test.go +++ b/cloud/deployment/deployment_test.go @@ -1372,6 +1372,26 @@ func TestCreate(t *testing.T) { assert.NoError(t, err) }) t.Run("list workspace failure", func(t *testing.T) { + mockClient.On("GetDeploymentConfigWithOrganization", mock.Anything).Return(astro.DeploymentConfig{ + Components: astro.Components{ + Scheduler: astro.SchedulerConfig{ + AU: astro.AuConfig{ + Default: 5, + Limit: 24, + }, + Replicas: astro.ReplicasConfig{ + Default: 1, + Minimum: 1, + Limit: 4, + }, + }, + }, + RuntimeReleases: []astro.RuntimeRelease{ + { + Version: "4.2.5", + }, + }, + }, nil).Times(1) mockClient.On("GetDeploymentConfig").Return(astro.DeploymentConfig{ Components: astro.Components{ Scheduler: astro.SchedulerConfig{ @@ -1391,13 +1411,33 @@ func TestCreate(t *testing.T) { Version: "4.2.5", }, }, - }, nil).Times(2) + }, nil).Times(1) mockCoreClient.On("ListWorkspacesWithResponse", mock.Anything, mock.Anything, mock.Anything).Return(nil, errMock).Once() err := Create("", ws, "test-desc", csID, "4.2.5", dagDeploy, CeleryExecutor, "", "", "", "", "", 10, 3, mockClient, mockCoreClient, false, &disableCiCdEnforcement) assert.ErrorIs(t, err, errMock) mockClient.AssertExpectations(t) }) t.Run("invalid workspace failure", func(t *testing.T) { + mockClient.On("GetDeploymentConfigWithOrganization", mock.Anything).Return(astro.DeploymentConfig{ + Components: astro.Components{ + Scheduler: astro.SchedulerConfig{ + AU: astro.AuConfig{ + Default: 5, + Limit: 24, + }, + Replicas: astro.ReplicasConfig{ + Default: 1, + Minimum: 1, + Limit: 4, + }, + }, + }, + RuntimeReleases: []astro.RuntimeRelease{ + { + Version: "4.2.5", + }, + }, + }, nil).Times(1) mockClient.On("GetDeploymentConfig").Return(astro.DeploymentConfig{ Components: astro.Components{ Scheduler: astro.SchedulerConfig{ @@ -1417,7 +1457,7 @@ func TestCreate(t *testing.T) { Version: "4.2.5", }, }, - }, nil).Times(2) + }, nil).Times(1) mockCoreClient.On("ListWorkspacesWithResponse", mock.Anything, mock.Anything, mock.Anything).Return(&ListWorkspacesResponseOK, nil).Once() err := Create("", "test-invalid-id", "test-desc", csID, "4.2.5", dagDeploy, CeleryExecutor, "", "", "", "", "", 10, 3, mockClient, mockCoreClient, false, &disableCiCdEnforcement) assert.Error(t, err) @@ -1431,6 +1471,26 @@ func TestCreate(t *testing.T) { ctx.SetContextKey("organization_product", "HOSTED") ctx.SetContextKey("organization", org) ctx.SetContextKey("organization_short_name", org) + mockClient.On("GetDeploymentConfigWithOrganization", mock.Anything).Return(astro.DeploymentConfig{ + Components: astro.Components{ + Scheduler: astro.SchedulerConfig{ + AU: astro.AuConfig{ + Default: 5, + Limit: 24, + }, + Replicas: astro.ReplicasConfig{ + Default: 1, + Minimum: 1, + Limit: 4, + }, + }, + }, + RuntimeReleases: []astro.RuntimeRelease{ + { + Version: "4.2.5", + }, + }, + }, nil).Times(1) mockClient.On("GetDeploymentConfig").Return(astro.DeploymentConfig{ Components: astro.Components{ Scheduler: astro.SchedulerConfig{ @@ -1450,7 +1510,7 @@ func TestCreate(t *testing.T) { Version: "4.2.5", }, }, - }, nil).Times(2) + }, nil).Times(1) mockCoreClient.On("ListWorkspacesWithResponse", mock.Anything, mock.Anything, mock.Anything).Return(&ListWorkspacesResponseOK, nil).Once() mockClient.On("CreateDeployment", &deploymentCreateInput).Return(astro.Deployment{ID: "test-id"}, nil).Once() mockClient.On("ListDeployments", org, ws).Return([]astro.Deployment{{ID: "test-id"}}, nil).Once() @@ -1476,6 +1536,26 @@ func TestCreate(t *testing.T) { mockCoreClient.AssertExpectations(t) }) t.Run("success with default config", func(t *testing.T) { + mockClient.On("GetDeploymentConfigWithOrganization", mock.Anything).Return(astro.DeploymentConfig{ + Components: astro.Components{ + Scheduler: astro.SchedulerConfig{ + AU: astro.AuConfig{ + Default: 5, + Limit: 24, + }, + Replicas: astro.ReplicasConfig{ + Default: 1, + Minimum: 1, + Limit: 4, + }, + }, + }, + RuntimeReleases: []astro.RuntimeRelease{ + { + Version: "4.2.5", + }, + }, + }, nil).Times(1) mockClient.On("GetDeploymentConfig").Return(astro.DeploymentConfig{ Components: astro.Components{ Scheduler: astro.SchedulerConfig{ @@ -1495,7 +1575,7 @@ func TestCreate(t *testing.T) { Version: "4.2.5", }, }, - }, nil).Times(2) + }, nil).Times(1) deploymentCreateInput := astro.CreateDeploymentInput{ WorkspaceID: ws, ClusterID: csID, From b0841be1d634c234c688b2c7528110718f562936 Mon Sep 17 00:00:00 2001 From: David Koenitzer Date: Thu, 19 Oct 2023 10:50:10 -0400 Subject: [PATCH 12/13] fix test --- cloud/deployment/deployment_test.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cloud/deployment/deployment_test.go b/cloud/deployment/deployment_test.go index 96da60fba..3db7778b0 100644 --- a/cloud/deployment/deployment_test.go +++ b/cloud/deployment/deployment_test.go @@ -1643,7 +1643,7 @@ func TestValidateResources(t *testing.T) { t.Run("invalid runtime version", func(t *testing.T) { mockClient := new(astro_mocks.Client) - mockClient.On("GetDeploymentConfig").Return(astro.DeploymentConfig{RuntimeReleases: []astro.RuntimeRelease{{Version: "4.2.5"}}}, nil).Once() + mockClient.On("GetDeploymentConfigWithOrganization", mock.Anything).Return(astro.DeploymentConfig{RuntimeReleases: []astro.RuntimeRelease{{Version: "4.2.5"}}}, nil).Once() resp, err := validateRuntimeVersion("4.2.4", mockClient) assert.NoError(t, err) From 8ddfd75980c24033c3c46066974b4afd00d56188 Mon Sep 17 00:00:00 2001 From: David Koenitzer Date: Thu, 19 Oct 2023 11:14:26 -0400 Subject: [PATCH 13/13] fix test --- astro-client/astro_test.go | 51 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 51 insertions(+) diff --git a/astro-client/astro_test.go b/astro-client/astro_test.go index 3b87b39ae..1cb545159 100644 --- a/astro-client/astro_test.go +++ b/astro-client/astro_test.go @@ -348,6 +348,57 @@ func TestGetDeploymentConfig(t *testing.T) { }) } +func TestGetDeploymentConfigWithOrganization(t *testing.T) { + testUtil.InitTestConfig(testUtil.CloudPlatform) + mockResponse := &Response{ + Data: ResponseData{ + GetDeploymentConfig: DeploymentConfig{ + AstronomerUnit: AstronomerUnit{CPU: 1, Memory: 1024}, + RuntimeReleases: []RuntimeRelease{ + { + Version: "4.2.5", + AirflowVersion: "2.2.5", + Channel: "stable", + ReleaseDate: "2020-06-25", + AirflowDatabaseMigration: true, + }, + }, + }, + }, + } + jsonResponse, err := json.Marshal(mockResponse) + assert.NoError(t, err) + + t.Run("success", func(t *testing.T) { + client := testUtil.NewTestClient(func(req *http.Request) *http.Response { + return &http.Response{ + StatusCode: 200, + Body: io.NopCloser(bytes.NewBuffer(jsonResponse)), + Header: make(http.Header), + } + }) + astroClient := NewAstroClient(client) + + deploymentConfig, err := astroClient.GetDeploymentConfigWithOrganization("test-org-id") + assert.NoError(t, err) + assert.Equal(t, deploymentConfig, mockResponse.Data.GetDeploymentConfig) + }) + + t.Run("error", func(t *testing.T) { + client := testUtil.NewTestClient(func(req *http.Request) *http.Response { + return &http.Response{ + StatusCode: 500, + Body: io.NopCloser(bytes.NewBufferString("Internal Server Error")), + Header: make(http.Header), + } + }) + astroClient := NewAstroClient(client) + + _, err := astroClient.GetDeploymentConfigWithOrganization("test-org-id") + assert.Contains(t, err.Error(), "Internal Server Error") + }) +} + func TestModifyDeploymentVariable(t *testing.T) { testUtil.InitTestConfig(testUtil.CloudPlatform) mockResponse := &Response{