FIx astro deploy <deployment-id> --image (#1464)

* astro deploy <deployment-id> --image does not work * update tests * update tests
astronomer · Nov 30, 2023 · 04848db · 04848db
1 parent aa8adb5
commit 04848db
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 3 deletions.
diff --git a/cloud/deploy/deploy.go b/cloud/deploy/deploy.go
@@ -564,11 +564,28 @@ func getImageName(cloudDomain, deploymentID, organizationID string, coreClient a
 	workspaceID := resp.JSON200.WorkspaceId
 	webserverURL := resp.JSON200.WebServerUrl
 	dagDeployEnabled := resp.JSON200.IsDagDeployEnabled
+	cicdEnforcement := resp.JSON200.ApiKeyOnlyDeployments
+	var desiredDagTarballVersion string
+	if resp.JSON200.DesiredDagTarballVersion != nil {
+		desiredDagTarballVersion = *resp.JSON200.DesiredDagTarballVersion
+	} else {
+		desiredDagTarballVersion = ""
+	}
 
 	// We use latest and keep this tag around after deployments to keep subsequent deploys quick
 	deployImage := airflow.ImageName(namespace, "latest")
 
-	return deploymentInfo{namespace: namespace, deployImage: deployImage, currentVersion: currentVersion, organizationID: organizationID, workspaceID: workspaceID, webserverURL: webserverURL, dagDeployEnabled: dagDeployEnabled}, nil
+	return deploymentInfo{
+		namespace:                namespace,
+		deployImage:              deployImage,
+		currentVersion:           currentVersion,
+		organizationID:           organizationID,
+		workspaceID:              workspaceID,
+		webserverURL:             webserverURL,
+		dagDeployEnabled:         dagDeployEnabled,
+		desiredDagTarballVersion: desiredDagTarballVersion,
+		cicdEnforcement:          cicdEnforcement,
+	}, nil
 }
 
 func buildImageWithoutDags(path string, imageHandler airflow.ImageHandler) error {

diff --git a/cloud/deploy/deploy_test.go b/cloud/deploy/deploy_test.go
@@ -63,7 +63,21 @@ var (
 			},
 		},
 	}
-	deploymentResponse = astrocore.GetDeploymentResponse{
+	DesiredDagTarballVersion = "desired-dag-tar-ball-version"
+	deploymentResponse       = astrocore.GetDeploymentResponse{
+		HTTPResponse: &http.Response{
+			StatusCode: 200,
+		},
+		JSON200: &astrocore.Deployment{
+			RuntimeVersion:           "4.2.5",
+			ReleaseName:              "test-name",
+			WorkspaceId:              ws,
+			WebServerUrl:             "test-url",
+			IsDagDeployEnabled:       false,
+			DesiredDagTarballVersion: &DesiredDagTarballVersion,
+		},
+	}
+	deploymentResponse2 = astrocore.GetDeploymentResponse{
 		HTTPResponse: &http.Response{
 			StatusCode: 200,
 		},
@@ -109,7 +123,7 @@ func TestDeployWithoutDagsDeploySuccess(t *testing.T) {
 	config.CFG.ShowWarnings.SetHomeString("false")
 	mockClient := new(astro_mocks.Client)
 
-	mockCoreClient.On("GetDeploymentWithResponse", mock.Anything, mock.Anything, mock.Anything).Return(&deploymentResponse, nil).Times(4)
+	mockCoreClient.On("GetDeploymentWithResponse", mock.Anything, mock.Anything, mock.Anything).Return(&deploymentResponse2, nil).Times(4)
 	mockCoreClient.On("ListDeploymentsWithResponse", mock.Anything, mock.Anything, mock.Anything).Return(&mockListDeploymentsResponse, nil).Times(1)
 	mockClient.On("ListDeployments", org, ws).Return([]astro.Deployment{{ID: "test-id", Workspace: astro.Workspace{ID: ws}}}, nil).Once()
 	mockCoreClient.On("GetDeploymentOptionsWithResponse", mock.Anything, mock.Anything, mock.Anything).Return(&getDeploymentOptionsResponse, nil).Times(5)