From 0e1d75e7f86faafcb18077745a3512cb6e135b75 Mon Sep 17 00:00:00 2001 From: Rujhan Arora Date: Thu, 12 Dec 2024 18:20:01 +0530 Subject: [PATCH 01/17] Added --image-name flag for software, similar to astri --- cmd/software/deploy.go | 4 +- cmd/software/deploy_test.go | 20 ++++-- software/deploy/deploy.go | 102 +++++++++++++++++++---------- software/deploy/deploy_test.go | 113 +++++++++++++++++++++++++++------ 4 files changed, 177 insertions(+), 62 deletions(-) diff --git a/cmd/software/deploy.go b/cmd/software/deploy.go index b305f2979..ed98bd413 100644 --- a/cmd/software/deploy.go +++ b/cmd/software/deploy.go @@ -27,6 +27,7 @@ var ( isDagOnlyDeploy bool description string isImageOnlyDeploy bool + imageName string ErrBothDagsOnlyAndImageOnlySet = errors.New("cannot use both --dags and --image together. Run 'astro deploy' to update both your image and dags") ) @@ -61,6 +62,7 @@ func NewDeployCmd() *cobra.Command { cmd.Flags().StringVar(&workspaceID, "workspace-id", "", "workspace assigned to deployment") cmd.Flags().StringVar(&description, "description", "", "Improve traceability by attaching a description to a code deploy. If you don't provide a description, the system automatically assigns a default description based on the deploy type.") cmd.Flags().BoolVarP(&isImageOnlyDeploy, "image", "", false, "Push only an image to your Astro Deployment. This only works for Dag-only, Git-sync-based and NFS-based deployments.") + cmd.Flags().StringVarP(&imageName, "image-name", "i", "", "Name of the custom image(present locally) to deploy") if !context.IsCloudContext() && houston.VerifyVersionMatch(houstonVersion, houston.VersionRestrictions{GTE: "0.34.0"}) { cmd.Flags().BoolVarP(&isDagOnlyDeploy, "dags", "d", false, "Push only DAGs to your Deployment") @@ -120,7 +122,7 @@ func deployAirflow(cmd *cobra.Command, args []string) error { } // Since we prompt the user to enter the deploymentID in come cases for DeployAirflowImage, reusing the same deploymentID for DagsOnlyDeploy - deploymentID, err = DeployAirflowImage(houstonClient, config.WorkingPath, deploymentID, ws, byoRegistryDomain, ignoreCacheDeploy, byoRegistryEnabled, forcePrompt, description, isImageOnlyDeploy) + deploymentID, err = DeployAirflowImage(houstonClient, config.WorkingPath, deploymentID, ws, byoRegistryDomain, ignoreCacheDeploy, byoRegistryEnabled, forcePrompt, description, isImageOnlyDeploy, imageName) if err != nil { return err } diff --git a/cmd/software/deploy_test.go b/cmd/software/deploy_test.go index 816ba88ec..eb2fc9b3a 100644 --- a/cmd/software/deploy_test.go +++ b/cmd/software/deploy_test.go @@ -27,7 +27,7 @@ func (s *Suite) TestDeploy() { EnsureProjectDir = func(cmd *cobra.Command, args []string) error { return nil } - DeployAirflowImage = func(houstonClient houston.ClientInterface, path, deploymentID, wsID, byoRegistryDomain string, ignoreCacheDeploy, byoRegistryEnabled, prompt bool, description string, isImageOnlyDeploy bool) (string, error) { + DeployAirflowImage = func(houstonClient houston.ClientInterface, path, deploymentID, wsID, byoRegistryDomain string, ignoreCacheDeploy, byoRegistryEnabled, prompt bool, description string, isImageOnlyDeploy bool, imageName string) (string, error) { if description == "" { return deploymentID, fmt.Errorf("description should not be empty") } @@ -52,7 +52,7 @@ func (s *Suite) TestDeploy() { s.NoError(err) // Test when the default description is used - DeployAirflowImage = func(houstonClient houston.ClientInterface, path, deploymentID, wsID, byoRegistryDomain string, ignoreCacheDeploy, byoRegistryEnabled, prompt bool, description string, isImageOnlyDeploy bool) (string, error) { + DeployAirflowImage = func(houstonClient houston.ClientInterface, path, deploymentID, wsID, byoRegistryDomain string, ignoreCacheDeploy, byoRegistryEnabled, prompt bool, description string, isImageOnlyDeploy bool, imageName string) (string, error) { expectedDesc := "Deployed via " if description != expectedDesc { return deploymentID, fmt.Errorf("expected description to be '%s', but got '%s'", expectedDesc, description) @@ -67,14 +67,14 @@ func (s *Suite) TestDeploy() { DagsOnlyDeploy = deploy.DagsOnlyDeploy s.Run("error should be returned for astro deploy, if DeployAirflowImage throws error", func() { - DeployAirflowImage = func(houstonClient houston.ClientInterface, path, deploymentID, wsID, byoRegistryDomain string, ignoreCacheDeploy, byoRegistryEnabled, prompt bool, description string, isImageOnlyDeploy bool) (string, error) { + DeployAirflowImage = func(houstonClient houston.ClientInterface, path, deploymentID, wsID, byoRegistryDomain string, ignoreCacheDeploy, byoRegistryEnabled, prompt bool, description string, isImageOnlyDeploy bool, imageName string) (string, error) { return deploymentID, deploy.ErrNoWorkspaceID } err := execDeployCmd([]string{"-f"}...) s.ErrorIs(err, deploy.ErrNoWorkspaceID) - DeployAirflowImage = func(houstonClient houston.ClientInterface, path, deploymentID, wsID, byoRegistryDomain string, ignoreCacheDeploy, byoRegistryEnabled, prompt bool, description string, isImageOnlyDeploy bool) (string, error) { + DeployAirflowImage = func(houstonClient houston.ClientInterface, path, deploymentID, wsID, byoRegistryDomain string, ignoreCacheDeploy, byoRegistryEnabled, prompt bool, description string, isImageOnlyDeploy bool, imageName string) (string, error) { return deploymentID, nil } }) @@ -104,7 +104,7 @@ func (s *Suite) TestDeploy() { }) s.Run("Test for the flag --image for image deployment", func() { - DeployAirflowImage = func(houstonClient houston.ClientInterface, path, deploymentID, wsID, byoRegistryDomain string, ignoreCacheDeploy, byoRegistryEnabled, prompt bool, description string, isImageOnlyDeploy bool) (string, error) { + DeployAirflowImage = func(houstonClient houston.ClientInterface, path, deploymentID, wsID, byoRegistryDomain string, ignoreCacheDeploy, byoRegistryEnabled, prompt bool, description string, isImageOnlyDeploy bool, imageName string) (string, error) { return deploymentID, deploy.ErrDeploymentTypeIncorrectForImageOnly } err := execDeployCmd([]string{"test-deployment-id", "--image", "--force"}...) @@ -112,7 +112,7 @@ func (s *Suite) TestDeploy() { }) s.Run("Test for the flag --image for dags-only deployment", func() { - DeployAirflowImage = func(houstonClient houston.ClientInterface, path, deploymentID, wsID, byoRegistryDomain string, ignoreCacheDeploy, byoRegistryEnabled, prompt bool, description string, isImageOnlyDeploy bool) (string, error) { + DeployAirflowImage = func(houstonClient houston.ClientInterface, path, deploymentID, wsID, byoRegistryDomain string, ignoreCacheDeploy, byoRegistryEnabled, prompt bool, description string, isImageOnlyDeploy bool, imageName string) (string, error) { return deploymentID, nil } // This function is not called since --image is passed @@ -137,4 +137,12 @@ func (s *Suite) TestDeploy() { s.ErrorIs(err, deploy.ErrBYORegistryDomainNotSet) DagsOnlyDeploy = deploy.DagsOnlyDeploy }) + + s.Run("Test for the flag --image-name for image deployment", func() { + DeployAirflowImage = func(houstonClient houston.ClientInterface, path, deploymentID, wsID, byoRegistryDomain string, ignoreCacheDeploy, byoRegistryEnabled, prompt bool, description string, isImageOnlyDeploy bool, imageName string) (string, error) { + return deploymentID, deploy.ErrDeploymentTypeIncorrectForImageOnly + } + err := execDeployCmd([]string{"test-deployment-id", "--image-name", "--force"}...) + s.ErrorIs(err, nil) + }) } diff --git a/software/deploy/deploy.go b/software/deploy/deploy.go index 501ed98e8..b37f0fc82 100644 --- a/software/deploy/deploy.go +++ b/software/deploy/deploy.go @@ -45,6 +45,8 @@ var ( ErrEmptyDagFolderUserCancelledOperation = errors.New("no DAGs found in the dags folder. User canceled the operation") ErrBYORegistryDomainNotSet = errors.New("Custom registry host is not set in config. It can be set at astronomer.houston.config.deployments.registry.protectedCustomRegistry.updateRegistry.host") //nolint ErrDeploymentTypeIncorrectForImageOnly = errors.New("--image only works for Dag-only, Git-sync-based and NFS-based deployments") + WarningInvalidImageNameMsg = "WARNING! The image in your Dockerfile '%s' is not based on Astro Runtime and is not supported. Change your Dockerfile with an image that pulls from 'quay.io/astronomer/astro-runtime' to proceed.\n" + ErrNoRuntimeLabelOnCustomImage = errors.New("the image should have label io.astronomer.docker.runtime.version") ) const ( @@ -58,9 +60,10 @@ const ( warningInvalidNameTag = "WARNING! You are about to push an image using the '%s' tag. This is not recommended.\nPlease use one of the following tags: %s.\nAre you sure you want to continue?" warningInvalidNameTagEmptyRecommendations = "WARNING! You are about to push an image using the '%s' tag. This is not recommended.\nAre you sure you want to continue?" - registryDomainPrefix = "registry." - runtimeImageLabel = "io.astronomer.docker.runtime.version" - airflowImageLabel = "io.astronomer.docker.airflow.version" + registryDomainPrefix = "registry." + runtimeImageLabel = "io.astronomer.docker.runtime.version" + airflowImageLabel = "io.astronomer.docker.airflow.version" + composeSkipImageBuildingPromptMsg = "Skipping building image since --image-name flag is used..." ) var tab = printutil.Table{ @@ -69,7 +72,7 @@ var tab = printutil.Table{ Header: []string{"#", "LABEL", "DEPLOYMENT NAME", "WORKSPACE", "DEPLOYMENT ID"}, } -func Airflow(houstonClient houston.ClientInterface, path, deploymentID, wsID, byoRegistryDomain string, ignoreCacheDeploy, byoRegistryEnabled, prompt bool, description string, isImageOnlyDeploy bool) (string, error) { +func Airflow(houstonClient houston.ClientInterface, path, deploymentID, wsID, byoRegistryDomain string, ignoreCacheDeploy, byoRegistryEnabled, prompt bool, description string, isImageOnlyDeploy bool, imageName string) (string, error) { deploymentID, deployments, err := getDeploymentIDForCurrentCommand(houstonClient, wsID, deploymentID, prompt) if err != nil { return deploymentID, err @@ -107,7 +110,7 @@ func Airflow(houstonClient houston.ClientInterface, path, deploymentID, wsID, by fmt.Printf(houstonDeploymentPrompt, releaseName) // Build the image to deploy - err = buildPushDockerImage(houstonClient, &c, deploymentInfo, releaseName, path, nextTag, cloudDomain, byoRegistryDomain, ignoreCacheDeploy, byoRegistryEnabled, description) + err = buildPushDockerImage(houstonClient, &c, deploymentInfo, releaseName, path, nextTag, cloudDomain, byoRegistryDomain, ignoreCacheDeploy, byoRegistryEnabled, description, imageName) if err != nil { return deploymentID, err } @@ -129,24 +132,7 @@ func deploymentExists(deploymentID string, deployments []houston.Deployment) boo return false } -func buildPushDockerImage(houstonClient houston.ClientInterface, c *config.Context, deploymentInfo *houston.Deployment, name, path, nextTag, cloudDomain, byoRegistryDomain string, ignoreCacheDeploy, byoRegistryEnabled bool, description string) error { - // Build our image - fmt.Println(imageBuildingPrompt) - - // parse dockerfile - cmds, err := docker.ParseFile(filepath.Join(path, dockerfile)) - if err != nil { - return fmt.Errorf("failed to parse dockerfile: %s: %w", filepath.Join(path, dockerfile), err) - } - - image, tag := docker.GetImageTagFromParsedFile(cmds) - if config.CFG.ShowWarnings.GetBool() && !validAirflowImageRepo(image) && !validRuntimeImageRepo(image) { - i, _ := input.Confirm(fmt.Sprintf(warningInvalidImageName, image)) - if !i { - fmt.Println("Canceling deploy...") - os.Exit(1) - } - } +func validateRuntimeVersion(houstonClient houston.ClientInterface, tag string, deploymentInfo *houston.Deployment) error { // Get valid image tags for platform using Deployment Info request deploymentConfig, err := houston.Call(houstonClient.GetDeploymentConfig)(nil) if err != nil { @@ -175,26 +161,72 @@ func buildPushDockerImage(houstonClient houston.ClientInterface, c *config.Conte os.Exit(1) } } - imageName := airflow.ImageName(name, "latest") + return nil +} +func buildPushDockerImage(houstonClient houston.ClientInterface, c *config.Context, deploymentInfo *houston.Deployment, name, path, nextTag, cloudDomain, byoRegistryDomain string, ignoreCacheDeploy, byoRegistryEnabled bool, description, customImageName string) error { + imageName := airflow.ImageName(name, "latest") imageHandler := imageHandlerInit(imageName) - if description != "" { deployLabels = append(deployLabels, "io.astronomer.deploy.revision.description="+description) } - buildConfig := types.ImageBuildConfig{ - Path: config.WorkingPath, - NoCache: ignoreCacheDeploy, - TargetPlatforms: deployImagePlatformSupport, - Output: true, - Labels: deployLabels, - } + var err error + if customImageName == "" { + // all these checks inside Dockerfile should happen only when no image-name is provided + // parse dockerfile + cmds, err := docker.ParseFile(filepath.Join(path, dockerfile)) + if err != nil { + return fmt.Errorf("failed to parse dockerfile: %s: %w", filepath.Join(path, dockerfile), err) + } - err = imageHandler.Build("", "", buildConfig) - if err != nil { - return err + image, tag := docker.GetImageTagFromParsedFile(cmds) + if config.CFG.ShowWarnings.GetBool() && !validAirflowImageRepo(image) && !validRuntimeImageRepo(image) { + i, _ := input.Confirm(fmt.Sprintf(warningInvalidImageName, image)) + if !i { + fmt.Println("Canceling deploy...") + os.Exit(1) + } + } + // Get valid image tags for platform using Deployment Info request + err = validateRuntimeVersion(houstonClient, tag, deploymentInfo) + if err != nil { + return err + } + // Build our image + fmt.Println(imageBuildingPrompt) + buildConfig := types.ImageBuildConfig{ + Path: config.WorkingPath, + NoCache: ignoreCacheDeploy, + TargetPlatforms: deployImagePlatformSupport, + Output: true, + Labels: deployLabels, + } + + err = imageHandler.Build("", "", buildConfig) + if err != nil { + return err + } + } else { + fmt.Println(composeSkipImageBuildingPromptMsg) + err = imageHandler.TagLocalImage(customImageName) + if err != nil { + return err + } + runtimeLabel, err := imageHandler.GetLabel("", airflow.RuntimeImageLabel) + if err != nil { + fmt.Println("unable get runtime version from image") + return err + } + if runtimeLabel == "" { + return ErrNoRuntimeLabelOnCustomImage + } + err = validateRuntimeVersion(houstonClient, runtimeLabel, deploymentInfo) + if err != nil { + return err + } } + var registry, remoteImage, token string if byoRegistryEnabled { registry = byoRegistryDomain diff --git a/software/deploy/deploy_test.go b/software/deploy/deploy_test.go index 3f1777f2e..b7e3d03ad 100644 --- a/software/deploy/deploy_test.go +++ b/software/deploy/deploy_test.go @@ -124,7 +124,7 @@ func (s *Suite) TestBuildPushDockerImageSuccessWithTagWarning() { houstonMock.On("GetRuntimeReleases", "").Return(houston.RuntimeReleases{}, nil) houstonMock.On("GetDeploymentConfig", nil).Return(mockedDeploymentConfig, nil) - err := buildPushDockerImage(houstonMock, &config.Context{}, mockDeployment, "test", "./testfiles/", "test", "test", "", false, false, description) + err := buildPushDockerImage(houstonMock, &config.Context{}, mockDeployment, "test", "./testfiles/", "test", "test", "", false, false, description, "") s.NoError(err) mockImageHandler.AssertExpectations(s.T()) houstonMock.AssertExpectations(s.T()) @@ -155,7 +155,7 @@ func (s *Suite) TestBuildPushDockerImageSuccessWithImageRepoWarning() { houstonMock.On("GetDeploymentConfig", nil).Return(mockedDeploymentConfig, nil) houstonMock.On("GetRuntimeReleases", "").Return(houston.RuntimeReleases{}, nil) - err := buildPushDockerImage(houstonMock, &config.Context{}, mockDeployment, "test", "./testfiles/", "test", "test", "", false, false, description) + err := buildPushDockerImage(houstonMock, &config.Context{}, mockDeployment, "test", "./testfiles/", "test", "test", "", false, false, description, "") s.NoError(err) mockImageHandler.AssertExpectations(s.T()) houstonMock.AssertExpectations(s.T()) @@ -202,7 +202,7 @@ func (s *Suite) TestBuildPushDockerImageSuccessWithBYORegistry() { houstonMock.On("GetRuntimeReleases", "").Return(houston.RuntimeReleases{}, nil) houstonMock.On("UpdateDeploymentImage", houston.UpdateDeploymentImageRequest{ReleaseName: "test", Image: "test.registry.io:test-test", AirflowVersion: "1.10.12", RuntimeVersion: ""}).Return(nil, nil) - err := buildPushDockerImage(houstonMock, &config.Context{}, mockDeployment, "test", "./testfiles/", "test", "test", "test.registry.io", false, true, description) + err := buildPushDockerImage(houstonMock, &config.Context{}, mockDeployment, "test", "./testfiles/", "test", "test", "test.registry.io", false, true, description, "") s.NoError(err) expectedLabel := deployRevisionDescriptionLabel + "=" + description @@ -235,7 +235,7 @@ func (s *Suite) TestBuildPushDockerImageSuccessWithBYORegistry() { } config.CFG.ShaAsTag.SetHomeString("true") defer config.CFG.ShaAsTag.SetHomeString("false") - err = buildPushDockerImage(houstonMock, &config.Context{}, mockDeployment, "test", "./testfiles/", "test", "test", "test.registry.io", false, true, description) + err = buildPushDockerImage(houstonMock, &config.Context{}, mockDeployment, "test", "./testfiles/", "test", "test", "test.registry.io", false, true, description, "") s.NoError(err) expectedLabel = deployRevisionDescriptionLabel + "=" + description assert.Contains(s.T(), capturedBuildConfig.Labels, expectedLabel) @@ -246,7 +246,7 @@ func (s *Suite) TestBuildPushDockerImageSuccessWithBYORegistry() { func (s *Suite) TestBuildPushDockerImageFailure() { // invalid dockerfile test dockerfile = "Dockerfile.invalid" - err := buildPushDockerImage(nil, &config.Context{}, mockDeployment, "test", "./testfiles/", "test", "test", "", false, false, description) + err := buildPushDockerImage(nil, &config.Context{}, mockDeployment, "test", "./testfiles/", "test", "test", "", false, false, description, "") s.EqualError(err, "failed to parse dockerfile: testfiles/Dockerfile.invalid: when using JSON array syntax, arrays must be comprised of strings only") dockerfile = "Dockerfile" @@ -262,7 +262,7 @@ func (s *Suite) TestBuildPushDockerImageFailure() { houstonMock.On("GetDeploymentConfig", nil).Return(nil, errMockHouston).Once() houstonMock.On("GetRuntimeReleases", "").Return(houston.RuntimeReleases{}, nil) // houston GetDeploymentConfig call failure - err = buildPushDockerImage(houstonMock, &config.Context{}, mockDeployment, "test", "./testfiles/", "test", "test", "", false, false, description) + err = buildPushDockerImage(houstonMock, &config.Context{}, mockDeployment, "test", "./testfiles/", "test", "test", "", false, false, description, "") s.Error(err, errMockHouston) houstonMock.On("GetDeploymentConfig", nil).Return(mockedDeploymentConfig, nil).Twice() @@ -274,7 +274,7 @@ func (s *Suite) TestBuildPushDockerImageFailure() { } // build error test case - err = buildPushDockerImage(houstonMock, &config.Context{}, mockDeployment, "test", "./testfiles/", "test", "test", "", false, false, description) + err = buildPushDockerImage(houstonMock, &config.Context{}, mockDeployment, "test", "./testfiles/", "test", "test", "", false, false, description, "") s.Error(err, errSomeContainerIssue.Error()) mockImageHandler.AssertExpectations(s.T()) @@ -286,7 +286,7 @@ func (s *Suite) TestBuildPushDockerImageFailure() { } // push error test case - err = buildPushDockerImage(houstonMock, &config.Context{}, mockDeployment, "test", "./testfiles/", "test", "test", "", false, false, description) + err = buildPushDockerImage(houstonMock, &config.Context{}, mockDeployment, "test", "./testfiles/", "test", "test", "", false, false, description, "") s.Error(err, errSomeContainerIssue.Error()) mockImageHandler.AssertExpectations(s.T()) houstonMock.AssertExpectations(s.T()) @@ -323,14 +323,14 @@ func (s *Suite) TestGetAirflowUILinkFailure() { func (s *Suite) TestAirflowFailure() { // No workspace ID test case - _, err := Airflow(nil, "", "", "", "", false, false, false, description, false) + _, err := Airflow(nil, "", "", "", "", false, false, false, description, false, "") s.ErrorIs(err, ErrNoWorkspaceID) // houston GetWorkspace failure case houstonMock := new(houston_mocks.ClientInterface) houstonMock.On("GetWorkspace", mock.Anything).Return(nil, errMockHouston).Once() - _, err = Airflow(houstonMock, "", "", "test-workspace-id", "", false, false, false, description, false) + _, err = Airflow(houstonMock, "", "", "test-workspace-id", "", false, false, false, description, false, "") s.ErrorIs(err, errMockHouston) houstonMock.AssertExpectations(s.T()) @@ -338,7 +338,7 @@ func (s *Suite) TestAirflowFailure() { houstonMock.On("GetWorkspace", mock.Anything).Return(&houston.Workspace{}, nil) houstonMock.On("ListDeployments", mock.Anything).Return(nil, errMockHouston).Once() - _, err = Airflow(houstonMock, "", "", "test-workspace-id", "", false, false, false, description, false) + _, err = Airflow(houstonMock, "", "", "test-workspace-id", "", false, false, false, description, false, "") s.ErrorIs(err, errMockHouston) houstonMock.AssertExpectations(s.T()) @@ -352,35 +352,35 @@ func (s *Suite) TestAirflowFailure() { // config GetCurrentContext failure case config.ResetCurrentContext() - _, err = Airflow(houstonMock, "", "", "test-workspace-id", "", false, false, false, description, false) + _, err = Airflow(houstonMock, "", "", "test-workspace-id", "", false, false, false, description, false, "") s.EqualError(err, "no context set, have you authenticated to Astro or Astronomer Software? Run astro login and try again") context.Switch("localhost") // Invalid deployment name case - _, err = Airflow(houstonMock, "", "test-deployment-id", "test-workspace-id", "", false, false, false, description, false) + _, err = Airflow(houstonMock, "", "test-deployment-id", "test-workspace-id", "", false, false, false, description, false, "") s.ErrorIs(err, errInvalidDeploymentID) // No deployment in the current workspace case - _, err = Airflow(houstonMock, "", "", "test-workspace-id", "", false, false, false, description, false) + _, err = Airflow(houstonMock, "", "", "test-workspace-id", "", false, false, false, description, false, "") s.ErrorIs(err, errDeploymentNotFound) houstonMock.AssertExpectations(s.T()) // Invalid deployment selection case houstonMock.On("ListDeployments", mock.Anything).Return([]houston.Deployment{{ID: "test-deployment-id"}}, nil) - _, err = Airflow(houstonMock, "", "", "test-workspace-id", "", false, false, false, description, false) + _, err = Airflow(houstonMock, "", "", "test-workspace-id", "", false, false, false, description, false, "") s.ErrorIs(err, errInvalidDeploymentSelected) // return error When houston get deployment throws an error houstonMock.On("ListDeployments", mock.Anything).Return([]houston.Deployment{{ID: "test-deployment-id"}}, nil) houstonMock.On("GetDeployment", mock.Anything).Return(nil, errMockHouston).Once() - _, err = Airflow(houstonMock, "", "test-deployment-id", "test-workspace-id", "", false, false, false, description, false) + _, err = Airflow(houstonMock, "", "test-deployment-id", "test-workspace-id", "", false, false, false, description, false, "") s.Equal(err.Error(), "failed to get deployment info: "+errMockHouston.Error()) // buildPushDockerImage failure case houstonMock.On("GetDeployment", "test-deployment-id").Return(&houston.Deployment{}, nil) dockerfile = "Dockerfile.invalid" - _, err = Airflow(houstonMock, "./testfiles/", "test-deployment-id", "test-workspace-id", "", false, false, false, description, false) + _, err = Airflow(houstonMock, "./testfiles/", "test-deployment-id", "test-workspace-id", "", false, false, false, description, false, "") dockerfile = "Dockerfile" s.Error(err) s.Contains(err.Error(), "failed to parse dockerfile") @@ -413,7 +413,7 @@ func (s *Suite) TestAirflowSuccess() { houstonMock.On("GetDeployment", mock.Anything).Return(&houston.Deployment{}, nil).Once() houstonMock.On("GetRuntimeReleases", "").Return(mockRuntimeReleases, nil) - _, err := Airflow(houstonMock, "./testfiles/", "test-deployment-id", "test-workspace-id", "", false, false, false, description, false) + _, err := Airflow(houstonMock, "./testfiles/", "test-deployment-id", "test-workspace-id", "", false, false, false, description, false, "") s.NoError(err) houstonMock.AssertExpectations(s.T()) } @@ -452,11 +452,84 @@ func (s *Suite) TestAirflowSuccessForImageOnly() { houstonMock.On("GetDeployment", mock.Anything).Return(deployment, nil).Once() houstonMock.On("GetRuntimeReleases", "").Return(mockRuntimeReleases, nil) - _, err := Airflow(houstonMock, "./testfiles/", "test-deployment-id", "test-workspace-id", "", false, false, false, description, true) + _, err := Airflow(houstonMock, "./testfiles/", "test-deployment-id", "test-workspace-id", "", false, false, false, description, true, "") s.NoError(err) houstonMock.AssertExpectations(s.T()) } +func (s *Suite) TestAirflowSuccessForImageName() { + fs := afero.NewMemMapFs() + configYaml := testUtil.NewTestConfig("localhost") + afero.WriteFile(fs, config.HomeConfigFile, configYaml, 0o777) + config.InitConfig(fs) + + customImageName := "test-image-name" + mockImageHandler := new(mocks.ImageHandler) + imageHandlerInit = func(image string) airflow.ImageHandler { + mockImageHandler.On("Push", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil) + mockImageHandler.On("TagLocalImage", customImageName).Return(nil) + mockImageHandler.On("GetLabel", "", airflow.RuntimeImageLabel).Return("test", nil) + return mockImageHandler + } + + mockedDeploymentConfig := &houston.DeploymentConfig{ + AirflowImages: mockAirflowImageList, + } + mockRuntimeReleases := houston.RuntimeReleases{ + houston.RuntimeRelease{Version: "4.2.4", AirflowVersion: "2.2.5"}, + houston.RuntimeRelease{Version: "4.2.5", AirflowVersion: "2.2.5"}, + } + houstonMock := new(houston_mocks.ClientInterface) + houstonMock.On("GetWorkspace", mock.Anything).Return(&houston.Workspace{}, nil).Once() + houstonMock.On("ListDeployments", mock.Anything).Return([]houston.Deployment{{ID: "test-deployment-id"}}, nil).Once() + houstonMock.On("GetDeploymentConfig", nil).Return(mockedDeploymentConfig, nil).Once() + dagDeployment := &houston.DagDeploymentConfig{ + Type: "dag-only", + } + deployment := &houston.Deployment{ + DagDeployment: *dagDeployment, + } + + houstonMock.On("GetDeployment", mock.Anything).Return(deployment, nil).Once() + houstonMock.On("GetRuntimeReleases", "").Return(mockRuntimeReleases, nil) + + _, err := Airflow(houstonMock, "./testfiles/", "test-deployment-id", "test-workspace-id", "", false, false, false, description, true, customImageName) + s.NoError(err) + houstonMock.AssertExpectations(s.T()) +} + +func (s *Suite) TestAirflowFailForImageNameWhenImageHasNoRuntimeLabel() { + fs := afero.NewMemMapFs() + configYaml := testUtil.NewTestConfig("localhost") + afero.WriteFile(fs, config.HomeConfigFile, configYaml, 0o777) + config.InitConfig(fs) + + customImageName := "test-image-name" + mockImageHandler := new(mocks.ImageHandler) + imageHandlerInit = func(image string) airflow.ImageHandler { + mockImageHandler.On("Push", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil) + mockImageHandler.On("TagLocalImage", customImageName).Return(nil) + mockImageHandler.On("GetLabel", "", airflow.RuntimeImageLabel).Return("", nil) + return mockImageHandler + } + + houstonMock := new(houston_mocks.ClientInterface) + houstonMock.On("GetWorkspace", mock.Anything).Return(&houston.Workspace{}, nil).Once() + houstonMock.On("ListDeployments", mock.Anything).Return([]houston.Deployment{{ID: "test-deployment-id"}}, nil).Once() + dagDeployment := &houston.DagDeploymentConfig{ + Type: "dag-only", + } + deployment := &houston.Deployment{ + DagDeployment: *dagDeployment, + } + + houstonMock.On("GetDeployment", mock.Anything).Return(deployment, nil).Once() + + _, err := Airflow(houstonMock, "./testfiles/", "test-deployment-id", "test-workspace-id", "", false, false, false, description, true, customImageName) + s.Error(err, ErrNoRuntimeLabelOnCustomImage) + houstonMock.AssertExpectations(s.T()) +} + func (s *Suite) TestAirflowFailureForImageOnly() { fs := afero.NewMemMapFs() configYaml := testUtil.NewTestConfig("localhost") @@ -481,7 +554,7 @@ func (s *Suite) TestAirflowFailureForImageOnly() { houstonMock.On("GetDeployment", mock.Anything).Return(deployment, nil).Once() - _, err := Airflow(houstonMock, "./testfiles/", "test-deployment-id", "test-workspace-id", "", false, false, false, description, true) + _, err := Airflow(houstonMock, "./testfiles/", "test-deployment-id", "test-workspace-id", "", false, false, false, description, true, "") s.Error(err, ErrDeploymentTypeIncorrectForImageOnly) houstonMock.AssertExpectations(s.T()) } From 6fcb6d0fd9f6d06011b83110cca77d51805412a2 Mon Sep 17 00:00:00 2001 From: Rujhan Arora Date: Thu, 12 Dec 2024 18:25:41 +0530 Subject: [PATCH 02/17] Fixed test --- cmd/software/deploy_test.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cmd/software/deploy_test.go b/cmd/software/deploy_test.go index eb2fc9b3a..770e08c93 100644 --- a/cmd/software/deploy_test.go +++ b/cmd/software/deploy_test.go @@ -140,7 +140,7 @@ func (s *Suite) TestDeploy() { s.Run("Test for the flag --image-name for image deployment", func() { DeployAirflowImage = func(houstonClient houston.ClientInterface, path, deploymentID, wsID, byoRegistryDomain string, ignoreCacheDeploy, byoRegistryEnabled, prompt bool, description string, isImageOnlyDeploy bool, imageName string) (string, error) { - return deploymentID, deploy.ErrDeploymentTypeIncorrectForImageOnly + return deploymentID, nil } err := execDeployCmd([]string{"test-deployment-id", "--image-name", "--force"}...) s.ErrorIs(err, nil) From 293253b57530f6ad2aa282ad799155e4a71dd7b4 Mon Sep 17 00:00:00 2001 From: Rujhan Arora Date: Thu, 12 Dec 2024 18:30:16 +0530 Subject: [PATCH 03/17] Fixed test --- cmd/software/deploy_test.go | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/cmd/software/deploy_test.go b/cmd/software/deploy_test.go index 770e08c93..183c674c3 100644 --- a/cmd/software/deploy_test.go +++ b/cmd/software/deploy_test.go @@ -123,6 +123,14 @@ func (s *Suite) TestDeploy() { s.ErrorIs(err, nil) }) + s.Run("Test for the flag --image-name for image deployment", func() { + DeployAirflowImage = func(houstonClient houston.ClientInterface, path, deploymentID, wsID, byoRegistryDomain string, ignoreCacheDeploy, byoRegistryEnabled, prompt bool, description string, isImageOnlyDeploy bool, imageName string) (string, error) { + return deploymentID, nil + } + err := execDeployCmd([]string{"test-deployment-id", "--image-name", "--force"}...) + s.ErrorIs(err, nil) + }) + s.Run("error should be returned if BYORegistryEnabled is true but BYORegistryDomain is empty", func() { appConfig = &houston.AppConfig{ BYORegistryDomain: "", @@ -137,12 +145,4 @@ func (s *Suite) TestDeploy() { s.ErrorIs(err, deploy.ErrBYORegistryDomainNotSet) DagsOnlyDeploy = deploy.DagsOnlyDeploy }) - - s.Run("Test for the flag --image-name for image deployment", func() { - DeployAirflowImage = func(houstonClient houston.ClientInterface, path, deploymentID, wsID, byoRegistryDomain string, ignoreCacheDeploy, byoRegistryEnabled, prompt bool, description string, isImageOnlyDeploy bool, imageName string) (string, error) { - return deploymentID, nil - } - err := execDeployCmd([]string{"test-deployment-id", "--image-name", "--force"}...) - s.ErrorIs(err, nil) - }) } From 67da2f76bd68488524e6763671867a128154d826 Mon Sep 17 00:00:00 2001 From: Rujhan Arora Date: Thu, 12 Dec 2024 18:38:50 +0530 Subject: [PATCH 04/17] Fixed test --- cmd/software/deploy_test.go | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/cmd/software/deploy_test.go b/cmd/software/deploy_test.go index 183c674c3..9d7890cab 100644 --- a/cmd/software/deploy_test.go +++ b/cmd/software/deploy_test.go @@ -123,11 +123,11 @@ func (s *Suite) TestDeploy() { s.ErrorIs(err, nil) }) - s.Run("Test for the flag --image-name for image deployment", func() { + s.Run("Test for the flag --image-name", func() { DeployAirflowImage = func(houstonClient houston.ClientInterface, path, deploymentID, wsID, byoRegistryDomain string, ignoreCacheDeploy, byoRegistryEnabled, prompt bool, description string, isImageOnlyDeploy bool, imageName string) (string, error) { return deploymentID, nil } - err := execDeployCmd([]string{"test-deployment-id", "--image-name", "--force"}...) + err := execDeployCmd([]string{"test-deployment-id", "--image-name", "--force", "--workspace-id=" + mockWorkspace.ID}...) s.ErrorIs(err, nil) }) From 4d75ce875ffd42e640619b351b55b4208c5829ce Mon Sep 17 00:00:00 2001 From: Rujhan Arora Date: Thu, 12 Dec 2024 18:51:37 +0530 Subject: [PATCH 05/17] Fixed test --- cmd/software/deploy_test.go | 3 +++ 1 file changed, 3 insertions(+) diff --git a/cmd/software/deploy_test.go b/cmd/software/deploy_test.go index 9d7890cab..8c4edc7f3 100644 --- a/cmd/software/deploy_test.go +++ b/cmd/software/deploy_test.go @@ -127,6 +127,9 @@ func (s *Suite) TestDeploy() { DeployAirflowImage = func(houstonClient houston.ClientInterface, path, deploymentID, wsID, byoRegistryDomain string, ignoreCacheDeploy, byoRegistryEnabled, prompt bool, description string, isImageOnlyDeploy bool, imageName string) (string, error) { return deploymentID, nil } + DagsOnlyDeploy = func(houstonClient houston.ClientInterface, appConfig *houston.AppConfig, wsID, deploymentID, dagsParentPath string, dagDeployURL *string, cleanUpFiles bool, description string) error { + return nil + } err := execDeployCmd([]string{"test-deployment-id", "--image-name", "--force", "--workspace-id=" + mockWorkspace.ID}...) s.ErrorIs(err, nil) }) From e4b32edff1617f957cae4bf98b2089c2e308f8b9 Mon Sep 17 00:00:00 2001 From: Rujhan Arora Date: Fri, 13 Dec 2024 12:27:57 +0530 Subject: [PATCH 06/17] 1. Split buildPushDockerImage into buildDockerImage(buildDockerImageForWorkingDir and buildDockerImageForCustomImage) and pushDockerImage. 2. Increased test coverage --- cmd/software/deploy_test.go | 7 +- software/deploy/deploy.go | 140 ++++++++++++++++++--------------- software/deploy/deploy_test.go | 5 +- 3 files changed, 85 insertions(+), 67 deletions(-) diff --git a/cmd/software/deploy_test.go b/cmd/software/deploy_test.go index 8c4edc7f3..2ae169904 100644 --- a/cmd/software/deploy_test.go +++ b/cmd/software/deploy_test.go @@ -124,14 +124,19 @@ func (s *Suite) TestDeploy() { }) s.Run("Test for the flag --image-name", func() { + var capturedImageName string DeployAirflowImage = func(houstonClient houston.ClientInterface, path, deploymentID, wsID, byoRegistryDomain string, ignoreCacheDeploy, byoRegistryEnabled, prompt bool, description string, isImageOnlyDeploy bool, imageName string) (string, error) { + capturedImageName = imageName // Capture the imageName return deploymentID, nil } DagsOnlyDeploy = func(houstonClient houston.ClientInterface, appConfig *houston.AppConfig, wsID, deploymentID, dagsParentPath string, dagDeployURL *string, cleanUpFiles bool, description string) error { return nil } - err := execDeployCmd([]string{"test-deployment-id", "--image-name", "--force", "--workspace-id=" + mockWorkspace.ID}...) + testImageName := "test-image-name" // Set the expected image name + err := execDeployCmd([]string{"test-deployment-id", "--image-name=" + testImageName, "--force", "--workspace-id=" + mockWorkspace.ID}...) + s.ErrorIs(err, nil) + s.Equal(testImageName, capturedImageName, "The imageName passed to DeployAirflowImage is incorrect") }) s.Run("error should be returned if BYORegistryEnabled is true but BYORegistryDomain is empty", func() { diff --git a/software/deploy/deploy.go b/software/deploy/deploy.go index b37f0fc82..a2e3097e2 100644 --- a/software/deploy/deploy.go +++ b/software/deploy/deploy.go @@ -164,69 +164,7 @@ func validateRuntimeVersion(houstonClient houston.ClientInterface, tag string, d return nil } -func buildPushDockerImage(houstonClient houston.ClientInterface, c *config.Context, deploymentInfo *houston.Deployment, name, path, nextTag, cloudDomain, byoRegistryDomain string, ignoreCacheDeploy, byoRegistryEnabled bool, description, customImageName string) error { - imageName := airflow.ImageName(name, "latest") - imageHandler := imageHandlerInit(imageName) - if description != "" { - deployLabels = append(deployLabels, "io.astronomer.deploy.revision.description="+description) - } - - var err error - if customImageName == "" { - // all these checks inside Dockerfile should happen only when no image-name is provided - // parse dockerfile - cmds, err := docker.ParseFile(filepath.Join(path, dockerfile)) - if err != nil { - return fmt.Errorf("failed to parse dockerfile: %s: %w", filepath.Join(path, dockerfile), err) - } - - image, tag := docker.GetImageTagFromParsedFile(cmds) - if config.CFG.ShowWarnings.GetBool() && !validAirflowImageRepo(image) && !validRuntimeImageRepo(image) { - i, _ := input.Confirm(fmt.Sprintf(warningInvalidImageName, image)) - if !i { - fmt.Println("Canceling deploy...") - os.Exit(1) - } - } - // Get valid image tags for platform using Deployment Info request - err = validateRuntimeVersion(houstonClient, tag, deploymentInfo) - if err != nil { - return err - } - // Build our image - fmt.Println(imageBuildingPrompt) - buildConfig := types.ImageBuildConfig{ - Path: config.WorkingPath, - NoCache: ignoreCacheDeploy, - TargetPlatforms: deployImagePlatformSupport, - Output: true, - Labels: deployLabels, - } - - err = imageHandler.Build("", "", buildConfig) - if err != nil { - return err - } - } else { - fmt.Println(composeSkipImageBuildingPromptMsg) - err = imageHandler.TagLocalImage(customImageName) - if err != nil { - return err - } - runtimeLabel, err := imageHandler.GetLabel("", airflow.RuntimeImageLabel) - if err != nil { - fmt.Println("unable get runtime version from image") - return err - } - if runtimeLabel == "" { - return ErrNoRuntimeLabelOnCustomImage - } - err = validateRuntimeVersion(houstonClient, runtimeLabel, deploymentInfo) - if err != nil { - return err - } - } - +func pushDockerImage(byoRegistryEnabled bool, byoRegistryDomain, name, nextTag, cloudDomain string, imageHandler airflow.ImageHandler, houstonClient houston.ClientInterface, c *config.Context) error { var registry, remoteImage, token string if byoRegistryEnabled { registry = byoRegistryDomain @@ -236,7 +174,7 @@ func buildPushDockerImage(houstonClient houston.ClientInterface, c *config.Conte remoteImage = fmt.Sprintf("%s/%s", registry, airflow.ImageName(name, nextTag)) token = c.Token } - err = imageHandler.Push(remoteImage, "", token) + err := imageHandler.Push(remoteImage, "", token) if err != nil { return err } @@ -255,10 +193,82 @@ func buildPushDockerImage(houstonClient houston.ClientInterface, c *config.Conte _, err := houston.Call(houstonClient.UpdateDeploymentImage)(req) return err } - return nil } +func buildDockerImageForCustomImage(imageHandler airflow.ImageHandler, customImageName string, deploymentInfo *houston.Deployment, houstonClient houston.ClientInterface) error { + fmt.Println(composeSkipImageBuildingPromptMsg) + err := imageHandler.TagLocalImage(customImageName) + if err != nil { + return err + } + runtimeLabel, err := imageHandler.GetLabel("", airflow.RuntimeImageLabel) + if err != nil { + fmt.Println("unable get runtime version from image") + return err + } + if runtimeLabel == "" { + return ErrNoRuntimeLabelOnCustomImage + } + err = validateRuntimeVersion(houstonClient, runtimeLabel, deploymentInfo) + return err +} + +func buildDockerImageForWorkingDir(path string, imageHandler airflow.ImageHandler, houstonClient houston.ClientInterface, deploymentInfo *houston.Deployment, ignoreCacheDeploy bool) error { + // all these checks inside Dockerfile should happen only when no image-name is provided + // parse dockerfile + cmds, err := docker.ParseFile(filepath.Join(path, dockerfile)) + if err != nil { + return fmt.Errorf("failed to parse dockerfile: %s: %w", filepath.Join(path, dockerfile), err) + } + + image, tag := docker.GetImageTagFromParsedFile(cmds) + if config.CFG.ShowWarnings.GetBool() && !validAirflowImageRepo(image) && !validRuntimeImageRepo(image) { + i, _ := input.Confirm(fmt.Sprintf(warningInvalidImageName, image)) + if !i { + fmt.Println("Canceling deploy...") + os.Exit(1) + } + } + // Get valid image tags for platform using Deployment Info request + err = validateRuntimeVersion(houstonClient, tag, deploymentInfo) + if err != nil { + return err + } + // Build our image + fmt.Println(imageBuildingPrompt) + buildConfig := types.ImageBuildConfig{ + Path: config.WorkingPath, + NoCache: ignoreCacheDeploy, + TargetPlatforms: deployImagePlatformSupport, + Output: true, + Labels: deployLabels, + } + + err = imageHandler.Build("", "", buildConfig) + return err +} + +func buildDockerImage(ignoreCacheDeploy bool, deploymentInfo *houston.Deployment, customImageName, path string, imageHandler airflow.ImageHandler, houstonClient houston.ClientInterface, c *config.Context) error { + if customImageName == "" { + return buildDockerImageForWorkingDir(path, imageHandler, houstonClient, deploymentInfo, ignoreCacheDeploy) + } + return buildDockerImageForCustomImage(imageHandler, customImageName, deploymentInfo, houstonClient) +} + +func buildPushDockerImage(houstonClient houston.ClientInterface, c *config.Context, deploymentInfo *houston.Deployment, name, path, nextTag, cloudDomain, byoRegistryDomain string, ignoreCacheDeploy, byoRegistryEnabled bool, description, customImageName string) error { + imageName := airflow.ImageName(name, "latest") + imageHandler := imageHandlerInit(imageName) + if description != "" { + deployLabels = append(deployLabels, "io.astronomer.deploy.revision.description="+description) + } + err := buildDockerImage(ignoreCacheDeploy, deploymentInfo, customImageName, path, imageHandler, houstonClient, c) + if err != nil { + return err + } + return pushDockerImage(byoRegistryEnabled, byoRegistryDomain, name, nextTag, cloudDomain, imageHandler, houstonClient, c) +} + func validAirflowImageRepo(image string) bool { validDockerfileBaseImages := map[string]bool{ "quay.io/astronomer/ap-airflow": true, diff --git a/software/deploy/deploy_test.go b/software/deploy/deploy_test.go index b7e3d03ad..c299b8843 100644 --- a/software/deploy/deploy_test.go +++ b/software/deploy/deploy_test.go @@ -455,6 +455,7 @@ func (s *Suite) TestAirflowSuccessForImageOnly() { _, err := Airflow(houstonMock, "./testfiles/", "test-deployment-id", "test-workspace-id", "", false, false, false, description, true, "") s.NoError(err) houstonMock.AssertExpectations(s.T()) + mockImageHandler.AssertExpectations(s.T()) } func (s *Suite) TestAirflowSuccessForImageName() { @@ -496,6 +497,7 @@ func (s *Suite) TestAirflowSuccessForImageName() { _, err := Airflow(houstonMock, "./testfiles/", "test-deployment-id", "test-workspace-id", "", false, false, false, description, true, customImageName) s.NoError(err) houstonMock.AssertExpectations(s.T()) + mockImageHandler.AssertExpectations(s.T()) } func (s *Suite) TestAirflowFailForImageNameWhenImageHasNoRuntimeLabel() { @@ -507,7 +509,6 @@ func (s *Suite) TestAirflowFailForImageNameWhenImageHasNoRuntimeLabel() { customImageName := "test-image-name" mockImageHandler := new(mocks.ImageHandler) imageHandlerInit = func(image string) airflow.ImageHandler { - mockImageHandler.On("Push", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil) mockImageHandler.On("TagLocalImage", customImageName).Return(nil) mockImageHandler.On("GetLabel", "", airflow.RuntimeImageLabel).Return("", nil) return mockImageHandler @@ -528,6 +529,7 @@ func (s *Suite) TestAirflowFailForImageNameWhenImageHasNoRuntimeLabel() { _, err := Airflow(houstonMock, "./testfiles/", "test-deployment-id", "test-workspace-id", "", false, false, false, description, true, customImageName) s.Error(err, ErrNoRuntimeLabelOnCustomImage) houstonMock.AssertExpectations(s.T()) + mockImageHandler.AssertExpectations(s.T()) } func (s *Suite) TestAirflowFailureForImageOnly() { @@ -557,6 +559,7 @@ func (s *Suite) TestAirflowFailureForImageOnly() { _, err := Airflow(houstonMock, "./testfiles/", "test-deployment-id", "test-workspace-id", "", false, false, false, description, true, "") s.Error(err, ErrDeploymentTypeIncorrectForImageOnly) houstonMock.AssertExpectations(s.T()) + mockImageHandler.AssertExpectations(s.T()) } func (s *Suite) TestDeployDagsOnlyFailure() { From 1bbd0e3d5fcfaaeebc339afa9532763250712fb8 Mon Sep 17 00:00:00 2001 From: Rujhan Arora Date: Fri, 13 Dec 2024 12:33:34 +0530 Subject: [PATCH 07/17] 1. Split buildPushDockerImage into buildDockerImage(buildDockerImageForWorkingDir and buildDockerImageForCustomImage) and pushDockerImage. 2. Increased test coverage --- software/deploy/deploy.go | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/software/deploy/deploy.go b/software/deploy/deploy.go index a2e3097e2..c46ab8fbd 100644 --- a/software/deploy/deploy.go +++ b/software/deploy/deploy.go @@ -249,7 +249,7 @@ func buildDockerImageForWorkingDir(path string, imageHandler airflow.ImageHandle return err } -func buildDockerImage(ignoreCacheDeploy bool, deploymentInfo *houston.Deployment, customImageName, path string, imageHandler airflow.ImageHandler, houstonClient houston.ClientInterface, c *config.Context) error { +func buildDockerImage(ignoreCacheDeploy bool, deploymentInfo *houston.Deployment, customImageName, path string, imageHandler airflow.ImageHandler, houstonClient houston.ClientInterface) error { if customImageName == "" { return buildDockerImageForWorkingDir(path, imageHandler, houstonClient, deploymentInfo, ignoreCacheDeploy) } @@ -262,7 +262,7 @@ func buildPushDockerImage(houstonClient houston.ClientInterface, c *config.Conte if description != "" { deployLabels = append(deployLabels, "io.astronomer.deploy.revision.description="+description) } - err := buildDockerImage(ignoreCacheDeploy, deploymentInfo, customImageName, path, imageHandler, houstonClient, c) + err := buildDockerImage(ignoreCacheDeploy, deploymentInfo, customImageName, path, imageHandler, houstonClient) if err != nil { return err } From cce4de038be9be66ec228750300786e139f9dad2 Mon Sep 17 00:00:00 2001 From: Rujhan Arora Date: Mon, 16 Dec 2024 15:25:12 +0530 Subject: [PATCH 08/17] Refactored code and added setup and teardown functionalities for --image-name tests --- software/deploy/deploy.go | 17 +++-- software/deploy/deploy_test.go | 113 ++++++++++++++++++--------------- 2 files changed, 70 insertions(+), 60 deletions(-) diff --git a/software/deploy/deploy.go b/software/deploy/deploy.go index c46ab8fbd..845027571 100644 --- a/software/deploy/deploy.go +++ b/software/deploy/deploy.go @@ -30,8 +30,6 @@ var ( gzipFile = fileutil.GzipFile getDeploymentIDForCurrentCommandVar = getDeploymentIDForCurrentCommand - - deployLabels = []string{"io.astronomer.skip.revision=true"} ) var ( @@ -214,7 +212,7 @@ func buildDockerImageForCustomImage(imageHandler airflow.ImageHandler, customIma return err } -func buildDockerImageForWorkingDir(path string, imageHandler airflow.ImageHandler, houstonClient houston.ClientInterface, deploymentInfo *houston.Deployment, ignoreCacheDeploy bool) error { +func buildDockerImageFromWorkingDir(path string, imageHandler airflow.ImageHandler, houstonClient houston.ClientInterface, deploymentInfo *houston.Deployment, ignoreCacheDeploy bool, description string) error { // all these checks inside Dockerfile should happen only when no image-name is provided // parse dockerfile cmds, err := docker.ParseFile(filepath.Join(path, dockerfile)) @@ -237,6 +235,10 @@ func buildDockerImageForWorkingDir(path string, imageHandler airflow.ImageHandle } // Build our image fmt.Println(imageBuildingPrompt) + deployLabels := []string{"io.astronomer.skip.revision=true"} + if description != "" { + deployLabels = append(deployLabels, "io.astronomer.deploy.revision.description="+description) + } buildConfig := types.ImageBuildConfig{ Path: config.WorkingPath, NoCache: ignoreCacheDeploy, @@ -249,9 +251,9 @@ func buildDockerImageForWorkingDir(path string, imageHandler airflow.ImageHandle return err } -func buildDockerImage(ignoreCacheDeploy bool, deploymentInfo *houston.Deployment, customImageName, path string, imageHandler airflow.ImageHandler, houstonClient houston.ClientInterface) error { +func buildDockerImage(ignoreCacheDeploy bool, deploymentInfo *houston.Deployment, customImageName, path string, imageHandler airflow.ImageHandler, houstonClient houston.ClientInterface, description string) error { if customImageName == "" { - return buildDockerImageForWorkingDir(path, imageHandler, houstonClient, deploymentInfo, ignoreCacheDeploy) + return buildDockerImageFromWorkingDir(path, imageHandler, houstonClient, deploymentInfo, ignoreCacheDeploy, description) } return buildDockerImageForCustomImage(imageHandler, customImageName, deploymentInfo, houstonClient) } @@ -259,10 +261,7 @@ func buildDockerImage(ignoreCacheDeploy bool, deploymentInfo *houston.Deployment func buildPushDockerImage(houstonClient houston.ClientInterface, c *config.Context, deploymentInfo *houston.Deployment, name, path, nextTag, cloudDomain, byoRegistryDomain string, ignoreCacheDeploy, byoRegistryEnabled bool, description, customImageName string) error { imageName := airflow.ImageName(name, "latest") imageHandler := imageHandlerInit(imageName) - if description != "" { - deployLabels = append(deployLabels, "io.astronomer.deploy.revision.description="+description) - } - err := buildDockerImage(ignoreCacheDeploy, deploymentInfo, customImageName, path, imageHandler, houstonClient) + err := buildDockerImage(ignoreCacheDeploy, deploymentInfo, customImageName, path, imageHandler, houstonClient, description) if err != nil { return err } diff --git a/software/deploy/deploy_test.go b/software/deploy/deploy_test.go index c299b8843..5b102f63b 100644 --- a/software/deploy/deploy_test.go +++ b/software/deploy/deploy_test.go @@ -65,6 +65,10 @@ var mockAirflowImageList = []houston.AirflowImage{ type Suite struct { suite.Suite + fsForDockerConfig afero.Fs + fsForLocalConfig afero.Fs + mockImageHandler *mocks.ImageHandler + houstonMock *houston_mocks.ClientInterface } func TestDeploy(t *testing.T) { @@ -99,6 +103,31 @@ func (s *Suite) TestValidRuntimeImageRepo() { s.False(validRuntimeImageRepo("personal-repo/ap-airflow")) } +func (s *Suite) SetupTest() { + // Common setup logic for the test suite + s.fsForLocalConfig = afero.NewMemMapFs() + afero.WriteFile(s.fsForLocalConfig, config.HomeConfigFile, testUtil.NewTestConfig("localhost"), 0o777) + + s.fsForDockerConfig = afero.NewMemMapFs() + afero.WriteFile(s.fsForLocalConfig, config.HomeConfigFile, testUtil.NewTestConfig("docker"), 0o777) + + s.mockImageHandler = new(mocks.ImageHandler) + imageHandlerInit = func(image string) airflow.ImageHandler { + return s.mockImageHandler + } + + s.houstonMock = new(houston_mocks.ClientInterface) +} + +func (s *Suite) TearDownSuite() { + // Cleanup logic, if any (e.g., clearing mocks) + s.mockImageHandler = nil + s.houstonMock = nil + s.fsForDockerConfig = nil + s.fsForLocalConfig = nil + imageHandlerInit = airflow.ImageHandlerInit +} + func (s *Suite) TestBuildPushDockerImageSuccessWithTagWarning() { fs := afero.NewMemMapFs() configYaml := testUtil.NewTestConfig("docker") @@ -459,18 +488,13 @@ func (s *Suite) TestAirflowSuccessForImageOnly() { } func (s *Suite) TestAirflowSuccessForImageName() { - fs := afero.NewMemMapFs() - configYaml := testUtil.NewTestConfig("localhost") - afero.WriteFile(fs, config.HomeConfigFile, configYaml, 0o777) - config.InitConfig(fs) - + config.InitConfig(s.fsForLocalConfig) customImageName := "test-image-name" - mockImageHandler := new(mocks.ImageHandler) imageHandlerInit = func(image string) airflow.ImageHandler { - mockImageHandler.On("Push", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil) - mockImageHandler.On("TagLocalImage", customImageName).Return(nil) - mockImageHandler.On("GetLabel", "", airflow.RuntimeImageLabel).Return("test", nil) - return mockImageHandler + s.mockImageHandler.On("Push", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil) + s.mockImageHandler.On("TagLocalImage", customImageName).Return(nil) + s.mockImageHandler.On("GetLabel", "", airflow.RuntimeImageLabel).Return("test", nil) + return s.mockImageHandler } mockedDeploymentConfig := &houston.DeploymentConfig{ @@ -480,10 +504,9 @@ func (s *Suite) TestAirflowSuccessForImageName() { houston.RuntimeRelease{Version: "4.2.4", AirflowVersion: "2.2.5"}, houston.RuntimeRelease{Version: "4.2.5", AirflowVersion: "2.2.5"}, } - houstonMock := new(houston_mocks.ClientInterface) - houstonMock.On("GetWorkspace", mock.Anything).Return(&houston.Workspace{}, nil).Once() - houstonMock.On("ListDeployments", mock.Anything).Return([]houston.Deployment{{ID: "test-deployment-id"}}, nil).Once() - houstonMock.On("GetDeploymentConfig", nil).Return(mockedDeploymentConfig, nil).Once() + s.houstonMock.On("GetWorkspace", mock.Anything).Return(&houston.Workspace{}, nil).Once() + s.houstonMock.On("ListDeployments", mock.Anything).Return([]houston.Deployment{{ID: "test-deployment-id"}}, nil).Once() + s.houstonMock.On("GetDeploymentConfig", nil).Return(mockedDeploymentConfig, nil).Once() dagDeployment := &houston.DagDeploymentConfig{ Type: "dag-only", } @@ -491,32 +514,26 @@ func (s *Suite) TestAirflowSuccessForImageName() { DagDeployment: *dagDeployment, } - houstonMock.On("GetDeployment", mock.Anything).Return(deployment, nil).Once() - houstonMock.On("GetRuntimeReleases", "").Return(mockRuntimeReleases, nil) + s.houstonMock.On("GetDeployment", mock.Anything).Return(deployment, nil).Once() + s.houstonMock.On("GetRuntimeReleases", "").Return(mockRuntimeReleases, nil) - _, err := Airflow(houstonMock, "./testfiles/", "test-deployment-id", "test-workspace-id", "", false, false, false, description, true, customImageName) + _, err := Airflow(s.houstonMock, "./testfiles/", "test-deployment-id", "test-workspace-id", "", false, false, false, description, true, customImageName) s.NoError(err) - houstonMock.AssertExpectations(s.T()) - mockImageHandler.AssertExpectations(s.T()) + s.houstonMock.AssertExpectations(s.T()) + s.mockImageHandler.AssertExpectations(s.T()) } func (s *Suite) TestAirflowFailForImageNameWhenImageHasNoRuntimeLabel() { - fs := afero.NewMemMapFs() - configYaml := testUtil.NewTestConfig("localhost") - afero.WriteFile(fs, config.HomeConfigFile, configYaml, 0o777) - config.InitConfig(fs) - + config.InitConfig(s.fsForLocalConfig) customImageName := "test-image-name" - mockImageHandler := new(mocks.ImageHandler) imageHandlerInit = func(image string) airflow.ImageHandler { - mockImageHandler.On("TagLocalImage", customImageName).Return(nil) - mockImageHandler.On("GetLabel", "", airflow.RuntimeImageLabel).Return("", nil) - return mockImageHandler + s.mockImageHandler.On("TagLocalImage", customImageName).Return(nil) + s.mockImageHandler.On("GetLabel", "", airflow.RuntimeImageLabel).Return("", nil) + return s.mockImageHandler } - houstonMock := new(houston_mocks.ClientInterface) - houstonMock.On("GetWorkspace", mock.Anything).Return(&houston.Workspace{}, nil).Once() - houstonMock.On("ListDeployments", mock.Anything).Return([]houston.Deployment{{ID: "test-deployment-id"}}, nil).Once() + s.houstonMock.On("GetWorkspace", mock.Anything).Return(&houston.Workspace{}, nil).Once() + s.houstonMock.On("ListDeployments", mock.Anything).Return([]houston.Deployment{{ID: "test-deployment-id"}}, nil).Once() dagDeployment := &houston.DagDeploymentConfig{ Type: "dag-only", } @@ -524,29 +541,23 @@ func (s *Suite) TestAirflowFailForImageNameWhenImageHasNoRuntimeLabel() { DagDeployment: *dagDeployment, } - houstonMock.On("GetDeployment", mock.Anything).Return(deployment, nil).Once() + s.houstonMock.On("GetDeployment", mock.Anything).Return(deployment, nil).Once() - _, err := Airflow(houstonMock, "./testfiles/", "test-deployment-id", "test-workspace-id", "", false, false, false, description, true, customImageName) + _, err := Airflow(s.houstonMock, "./testfiles/", "test-deployment-id", "test-workspace-id", "", false, false, false, description, true, customImageName) s.Error(err, ErrNoRuntimeLabelOnCustomImage) - houstonMock.AssertExpectations(s.T()) - mockImageHandler.AssertExpectations(s.T()) + s.houstonMock.AssertExpectations(s.T()) + s.mockImageHandler.AssertExpectations(s.T()) } func (s *Suite) TestAirflowFailureForImageOnly() { - fs := afero.NewMemMapFs() - configYaml := testUtil.NewTestConfig("localhost") - afero.WriteFile(fs, config.HomeConfigFile, configYaml, 0o777) - config.InitConfig(fs) - - mockImageHandler := new(mocks.ImageHandler) + config.InitConfig(s.fsForLocalConfig) imageHandlerInit = func(image string) airflow.ImageHandler { - mockImageHandler.On("Build", mock.Anything, mock.Anything, mock.Anything).Return(nil) - mockImageHandler.On("Push", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil) - return mockImageHandler + s.mockImageHandler.On("Build", mock.Anything, mock.Anything, mock.Anything).Return(nil) + s.mockImageHandler.On("Push", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil) + return s.mockImageHandler } - houstonMock := new(houston_mocks.ClientInterface) - houstonMock.On("GetWorkspace", mock.Anything).Return(&houston.Workspace{}, nil).Once() - houstonMock.On("ListDeployments", mock.Anything).Return([]houston.Deployment{{ID: "test-deployment-id"}}, nil).Once() + s.houstonMock.On("GetWorkspace", mock.Anything).Return(&houston.Workspace{}, nil).Once() + s.houstonMock.On("ListDeployments", mock.Anything).Return([]houston.Deployment{{ID: "test-deployment-id"}}, nil).Once() dagDeployment := &houston.DagDeploymentConfig{ Type: "image", } @@ -554,12 +565,12 @@ func (s *Suite) TestAirflowFailureForImageOnly() { DagDeployment: *dagDeployment, } - houstonMock.On("GetDeployment", mock.Anything).Return(deployment, nil).Once() + s.houstonMock.On("GetDeployment", mock.Anything).Return(deployment, nil).Once() - _, err := Airflow(houstonMock, "./testfiles/", "test-deployment-id", "test-workspace-id", "", false, false, false, description, true, "") + _, err := Airflow(s.houstonMock, "./testfiles/", "test-deployment-id", "test-workspace-id", "", false, false, false, description, true, "") s.Error(err, ErrDeploymentTypeIncorrectForImageOnly) - houstonMock.AssertExpectations(s.T()) - mockImageHandler.AssertExpectations(s.T()) + s.houstonMock.AssertExpectations(s.T()) + s.mockImageHandler.AssertExpectations(s.T()) } func (s *Suite) TestDeployDagsOnlyFailure() { From 2764940f6883dac2c253d3ac35c4cc7fb798ec1b Mon Sep 17 00:00:00 2001 From: Rujhan Arora Date: Mon, 16 Dec 2024 15:44:11 +0530 Subject: [PATCH 09/17] Replaced SetupTEst with SetupSuite --- software/deploy/deploy_test.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/software/deploy/deploy_test.go b/software/deploy/deploy_test.go index 5b102f63b..6ad60dcc0 100644 --- a/software/deploy/deploy_test.go +++ b/software/deploy/deploy_test.go @@ -103,7 +103,7 @@ func (s *Suite) TestValidRuntimeImageRepo() { s.False(validRuntimeImageRepo("personal-repo/ap-airflow")) } -func (s *Suite) SetupTest() { +func (s *Suite) SetupSuite() { // Common setup logic for the test suite s.fsForLocalConfig = afero.NewMemMapFs() afero.WriteFile(s.fsForLocalConfig, config.HomeConfigFile, testUtil.NewTestConfig("localhost"), 0o777) From c20a0c3ae58697c42a459087c0c7f567c416684e Mon Sep 17 00:00:00 2001 From: Rujhan Arora Date: Mon, 16 Dec 2024 15:48:33 +0530 Subject: [PATCH 10/17] Replaced SetupTEst with SetupSuite --- software/deploy/deploy_test.go | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/software/deploy/deploy_test.go b/software/deploy/deploy_test.go index 6ad60dcc0..0967a521f 100644 --- a/software/deploy/deploy_test.go +++ b/software/deploy/deploy_test.go @@ -110,12 +110,13 @@ func (s *Suite) SetupSuite() { s.fsForDockerConfig = afero.NewMemMapFs() afero.WriteFile(s.fsForLocalConfig, config.HomeConfigFile, testUtil.NewTestConfig("docker"), 0o777) +} +func (s *Suite) SetupTest() { s.mockImageHandler = new(mocks.ImageHandler) imageHandlerInit = func(image string) airflow.ImageHandler { return s.mockImageHandler } - s.houstonMock = new(houston_mocks.ClientInterface) } From dc21a6fd132fd50a6f8a2ceef6c9365f70fa76df Mon Sep 17 00:00:00 2001 From: Rujhan Arora Date: Tue, 17 Dec 2024 12:50:07 +0530 Subject: [PATCH 11/17] Moved houstonMock and mockImageHandler to teardown test and subtest functions for deploy_test.go --- software/deploy/deploy_test.go | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/software/deploy/deploy_test.go b/software/deploy/deploy_test.go index 0967a521f..9bf286b92 100644 --- a/software/deploy/deploy_test.go +++ b/software/deploy/deploy_test.go @@ -129,6 +129,16 @@ func (s *Suite) TearDownSuite() { imageHandlerInit = airflow.ImageHandlerInit } +func (s *Suite) TearDownSubTest() { + s.houstonMock.AssertExpectations(s.T()) + s.mockImageHandler.AssertExpectations(s.T()) +} + +func (s *Suite) TearDownTest() { + s.houstonMock.AssertExpectations(s.T()) + s.mockImageHandler.AssertExpectations(s.T()) +} + func (s *Suite) TestBuildPushDockerImageSuccessWithTagWarning() { fs := afero.NewMemMapFs() configYaml := testUtil.NewTestConfig("docker") @@ -520,8 +530,6 @@ func (s *Suite) TestAirflowSuccessForImageName() { _, err := Airflow(s.houstonMock, "./testfiles/", "test-deployment-id", "test-workspace-id", "", false, false, false, description, true, customImageName) s.NoError(err) - s.houstonMock.AssertExpectations(s.T()) - s.mockImageHandler.AssertExpectations(s.T()) } func (s *Suite) TestAirflowFailForImageNameWhenImageHasNoRuntimeLabel() { @@ -546,8 +554,6 @@ func (s *Suite) TestAirflowFailForImageNameWhenImageHasNoRuntimeLabel() { _, err := Airflow(s.houstonMock, "./testfiles/", "test-deployment-id", "test-workspace-id", "", false, false, false, description, true, customImageName) s.Error(err, ErrNoRuntimeLabelOnCustomImage) - s.houstonMock.AssertExpectations(s.T()) - s.mockImageHandler.AssertExpectations(s.T()) } func (s *Suite) TestAirflowFailureForImageOnly() { @@ -570,8 +576,6 @@ func (s *Suite) TestAirflowFailureForImageOnly() { _, err := Airflow(s.houstonMock, "./testfiles/", "test-deployment-id", "test-workspace-id", "", false, false, false, description, true, "") s.Error(err, ErrDeploymentTypeIncorrectForImageOnly) - s.houstonMock.AssertExpectations(s.T()) - s.mockImageHandler.AssertExpectations(s.T()) } func (s *Suite) TestDeployDagsOnlyFailure() { From fba9fd2c9851f694c9d121b7dca209b5f02f1069 Mon Sep 17 00:00:00 2001 From: Rujhan Arora Date: Fri, 13 Dec 2024 17:03:27 +0530 Subject: [PATCH 12/17] Introduction of astro deploy --image-name=image_name --remote --runtime-version=rt_version --- cmd/software/deploy.go | 23 +++++++--- cmd/software/deploy_test.go | 23 ++++++++++ software/deploy/deploy.go | 55 ++++++++++++++++++------ software/deploy/deploy_test.go | 78 ++++++++++++++++++++++++++++++++++ 4 files changed, 160 insertions(+), 19 deletions(-) diff --git a/cmd/software/deploy.go b/cmd/software/deploy.go index 7abbdaeb2..d4305db89 100644 --- a/cmd/software/deploy.go +++ b/cmd/software/deploy.go @@ -24,10 +24,13 @@ var ( EnsureProjectDir = utils.EnsureProjectDir DeployAirflowImage = deploy.Airflow DagsOnlyDeploy = deploy.DagsOnlyDeploy + UpdateDeploymentImage = deploy.UpdateDeploymentImage isDagOnlyDeploy bool description string isImageOnlyDeploy bool imageName string + runtimeVersionForImageName string + imagePresentOnRemote bool ErrBothDagsOnlyAndImageOnlySet = errors.New("cannot use both --dags and --image together. Run 'astro deploy' to update both your image and dags") ) @@ -62,7 +65,9 @@ func NewDeployCmd() *cobra.Command { cmd.Flags().StringVar(&workspaceID, "workspace-id", "", "workspace assigned to deployment") cmd.Flags().StringVar(&description, "description", "", "Improve traceability by attaching a description to a code deploy. If you don't provide a description, the system automatically assigns a default description based on the deploy type.") cmd.Flags().BoolVarP(&isImageOnlyDeploy, "image", "", false, "Push only an image to your Astro Deployment. This only works for Dag-only, Git-sync-based and NFS-based deployments.") - cmd.Flags().StringVarP(&imageName, "image-name", "i", "", "Name of the custom image(present locally) to deploy") + cmd.Flags().StringVarP(&imageName, "image-name", "i", "", "Name of the custom image(should be present locally unless --remote is specified) to deploy") + cmd.Flags().StringVar(&runtimeVersionForImageName, "runtime-version", "", "Runtime version of the image to deploy. Example - 12.1.1. Mandatory if --image-name --remote is provided") + cmd.Flags().BoolVarP(&imagePresentOnRemote, "remote", "", false, "Custom image is present on remote registry") if !context.IsCloudContext() && houston.VerifyVersionMatch(houstonVersion, houston.VersionRestrictions{GTE: "0.34.0"}) { cmd.Flags().BoolVarP(&isDagOnlyDeploy, "dags", "d", false, "Push only DAGs to your Deployment") @@ -121,11 +126,19 @@ func deployAirflow(cmd *cobra.Command, args []string) error { return DagsOnlyDeploy(houstonClient, appConfig, ws, deploymentID, config.WorkingPath, nil, true, description) } - // Since we prompt the user to enter the deploymentID in come cases for DeployAirflowImage, reusing the same deploymentID for DagsOnlyDeploy - deploymentID, err = DeployAirflowImage(houstonClient, config.WorkingPath, deploymentID, ws, byoRegistryDomain, ignoreCacheDeploy, byoRegistryEnabled, forcePrompt, description, isImageOnlyDeploy, imageName) - if err != nil { - return err + if imagePresentOnRemote && imageName != "" { + deploymentID, err = UpdateDeploymentImage(houstonClient, deploymentID, ws, runtimeVersionForImageName, imageName) + if err != nil { + return err + } + } else { + // Since we prompt the user to enter the deploymentID in come cases for DeployAirflowImage, reusing the same deploymentID for DagsOnlyDeploy + deploymentID, err = DeployAirflowImage(houstonClient, config.WorkingPath, deploymentID, ws, byoRegistryDomain, ignoreCacheDeploy, byoRegistryEnabled, forcePrompt, description, isImageOnlyDeploy, imageName) + if err != nil { + return err + } } + // Don't deploy dags even for dags-only deployments --image is passed if isImageOnlyDeploy { fmt.Println("Dags in the project will not be deployed since --image is passed.") diff --git a/cmd/software/deploy_test.go b/cmd/software/deploy_test.go index 2ae169904..0ad9ba426 100644 --- a/cmd/software/deploy_test.go +++ b/cmd/software/deploy_test.go @@ -139,6 +139,29 @@ func (s *Suite) TestDeploy() { s.Equal(testImageName, capturedImageName, "The imageName passed to DeployAirflowImage is incorrect") }) + s.Run("Test for the flag --image-name with --remote. Dags should be deployed but DeployAirflowImage shouldn't be called", func() { + DagsOnlyDeploy = func(houstonClient houston.ClientInterface, appConfig *houston.AppConfig, wsID, deploymentID, dagsParentPath string, dagDeployURL *string, cleanUpFiles bool, description string) error { + return nil + } + UpdateDeploymentImage = func(houstonClient houston.ClientInterface, deploymentID, wsID, runtimeVersion, imageName string) (string, error) { + return "", nil + } + testImageName := "test-image-name" // Set the expected image name + err := execDeployCmd([]string{"test-deployment-id", "--image-name=" + testImageName, "--force", "--remote", "--workspace-id=" + mockWorkspace.ID}...) + s.ErrorIs(err, nil) + UpdateDeploymentImage = deploy.UpdateDeploymentImage + }) + + s.Run("Test for the flag --image-name with --remote. Dags should not be deployed if UpdateDeploymentImage throws an error", func() { + UpdateDeploymentImage = func(houstonClient houston.ClientInterface, deploymentID, wsID, runtimeVersion, imageName string) (string, error) { + return "", errNoWorkspaceFound + } + testImageName := "test-image-name" // Set the expected image name + err := execDeployCmd([]string{"test-deployment-id", "--image-name=" + testImageName, "--force", "--remote", "--workspace-id=" + mockWorkspace.ID}...) + s.ErrorIs(err, errNoWorkspaceFound) + UpdateDeploymentImage = deploy.UpdateDeploymentImage + }) + s.Run("error should be returned if BYORegistryEnabled is true but BYORegistryDomain is empty", func() { appConfig = &houston.AppConfig{ BYORegistryDomain: "", diff --git a/software/deploy/deploy.go b/software/deploy/deploy.go index 845027571..2eb8a9251 100644 --- a/software/deploy/deploy.go +++ b/software/deploy/deploy.go @@ -33,18 +33,19 @@ var ( ) var ( - ErrNoWorkspaceID = errors.New("no workspace id provided") - errNoDomainSet = errors.New("no domain set, re-authenticate") - errInvalidDeploymentID = errors.New("please specify a valid deployment ID") - errDeploymentNotFound = errors.New("no airflow deployments found") - errInvalidDeploymentSelected = errors.New("invalid deployment selection\n") //nolint - ErrDagOnlyDeployDisabledInConfig = errors.New("to perform this operation, set both deployments.dagOnlyDeployment and deployments.configureDagDeployment to true in your Astronomer cluster") - ErrDagOnlyDeployNotEnabledForDeployment = errors.New("to perform this operation, first set the Deployment type to 'dag_deploy' via the UI or the API or the CLI") - ErrEmptyDagFolderUserCancelledOperation = errors.New("no DAGs found in the dags folder. User canceled the operation") - ErrBYORegistryDomainNotSet = errors.New("Custom registry host is not set in config. It can be set at astronomer.houston.config.deployments.registry.protectedCustomRegistry.updateRegistry.host") //nolint - ErrDeploymentTypeIncorrectForImageOnly = errors.New("--image only works for Dag-only, Git-sync-based and NFS-based deployments") - WarningInvalidImageNameMsg = "WARNING! The image in your Dockerfile '%s' is not based on Astro Runtime and is not supported. Change your Dockerfile with an image that pulls from 'quay.io/astronomer/astro-runtime' to proceed.\n" - ErrNoRuntimeLabelOnCustomImage = errors.New("the image should have label io.astronomer.docker.runtime.version") + ErrNoWorkspaceID = errors.New("no workspace id provided") + errNoDomainSet = errors.New("no domain set, re-authenticate") + errInvalidDeploymentID = errors.New("please specify a valid deployment ID") + errDeploymentNotFound = errors.New("no airflow deployments found") + errInvalidDeploymentSelected = errors.New("invalid deployment selection\n") //nolint + ErrDagOnlyDeployDisabledInConfig = errors.New("to perform this operation, set both deployments.dagOnlyDeployment and deployments.configureDagDeployment to true in your Astronomer cluster") + ErrDagOnlyDeployNotEnabledForDeployment = errors.New("to perform this operation, first set the Deployment type to 'dag_deploy' via the UI or the API or the CLI") + ErrEmptyDagFolderUserCancelledOperation = errors.New("no DAGs found in the dags folder. User canceled the operation") + ErrBYORegistryDomainNotSet = errors.New("Custom registry host is not set in config. It can be set at astronomer.houston.config.deployments.registry.protectedCustomRegistry.updateRegistry.host") //nolint + ErrDeploymentTypeIncorrectForImageOnly = errors.New("--image only works for Dag-only, Git-sync-based and NFS-based deployments") + WarningInvalidImageNameMsg = "WARNING! The image in your Dockerfile '%s' is not based on Astro Runtime and is not supported. Change your Dockerfile with an image that pulls from 'quay.io/astronomer/astro-runtime' to proceed.\n" + ErrNoRuntimeLabelOnCustomImage = errors.New("the image should have label io.astronomer.docker.runtime.version") + ErrRuntimeVersionNotPassedForRemoteImage = errors.New("if --image-name and --remote is passed, it's mandatory to pass --runtime-version") ) const ( @@ -162,6 +163,33 @@ func validateRuntimeVersion(houstonClient houston.ClientInterface, tag string, d return nil } +func updateDeploymentImageAPICall(houstonClient houston.ClientInterface, imageName, releaseName, airflowVersion, runtimeVersion string) (interface{}, error) { + req := houston.UpdateDeploymentImageRequest{ReleaseName: releaseName, Image: imageName, AirflowVersion: airflowVersion, RuntimeVersion: runtimeVersion} + return houston.Call(houstonClient.UpdateDeploymentImage)(req) +} + +func UpdateDeploymentImage(houstonClient houston.ClientInterface, deploymentID, wsID, runtimeVersion, imageName string) (string, error) { + if runtimeVersion == "" { + return "", ErrRuntimeVersionNotPassedForRemoteImage + } + deploymentIDForCurrentCmd, _, err := getDeploymentIDForCurrentCommandVar(houstonClient, wsID, deploymentID, deploymentID == "") + if err != nil { + return "", err + } + deploymentID = deploymentIDForCurrentCmd + if deploymentID == "" { + return "", errInvalidDeploymentID + } + deploymentInfo, err := houston.Call(houstonClient.GetDeployment)(deploymentID) + if err != nil { + return "", fmt.Errorf("failed to get deployment info: %w", err) + } + fmt.Println("Skipping building the image since --image-name flag is used...") + _, err = updateDeploymentImageAPICall(houstonClient, imageName, deploymentInfo.ReleaseName, "", runtimeVersion) + fmt.Println("Image successfully updated") + return deploymentIDForCurrentCmd, err +} + func pushDockerImage(byoRegistryEnabled bool, byoRegistryDomain, name, nextTag, cloudDomain string, imageHandler airflow.ImageHandler, houstonClient houston.ClientInterface, c *config.Context) error { var registry, remoteImage, token string if byoRegistryEnabled { @@ -187,8 +215,7 @@ func pushDockerImage(byoRegistryEnabled bool, byoRegistryDomain, name, nextTag, } runtimeVersion, _ := imageHandler.GetLabel("", runtimeImageLabel) airflowVersion, _ := imageHandler.GetLabel("", airflowImageLabel) - req := houston.UpdateDeploymentImageRequest{ReleaseName: name, Image: remoteImage, AirflowVersion: airflowVersion, RuntimeVersion: runtimeVersion} - _, err := houston.Call(houstonClient.UpdateDeploymentImage)(req) + _, err := updateDeploymentImageAPICall(houstonClient, remoteImage, name, airflowVersion, runtimeVersion) return err } return nil diff --git a/software/deploy/deploy_test.go b/software/deploy/deploy_test.go index 9bf286b92..07dfa328e 100644 --- a/software/deploy/deploy_test.go +++ b/software/deploy/deploy_test.go @@ -1033,3 +1033,81 @@ func (s *Suite) TestDeployDagsOnlyFailure() { s.True(os.IsNotExist(err)) }) } + +func (s *Suite) TestUpdateDeploymentImage() { + testUtil.InitTestConfig(testUtil.SoftwarePlatform) + deploymentID := "test-deployment-id" + wsID := "test-workspace-id" + runtimeVersion := "12.1.1" + imageName := "imageName" + releaseName := "releaseName" + + s.Run("When runtimeVersion is empty", func() { + houstonMock := new(houston_mocks.ClientInterface) + returnedDeploymentId, err := UpdateDeploymentImage(houstonMock, deploymentID, wsID, "", imageName) + s.ErrorIs(err, ErrRuntimeVersionNotPassedForRemoteImage) + houstonMock.AssertExpectations(s.T()) + s.Equal(returnedDeploymentId, "") + }) + + s.Run("When getDeploymentIDForCurrentCommandVar gives an error", func() { + getDeploymentIDForCurrentCommandVar = func(houstonClient houston.ClientInterface, wsID, deploymentID string, prompt bool) (string, []houston.Deployment, error) { + return deploymentID, nil, errDeploymentNotFound + } + houstonMock := new(houston_mocks.ClientInterface) + returnedDeploymentId, err := UpdateDeploymentImage(houstonMock, deploymentID, wsID, runtimeVersion, imageName) + s.ErrorIs(err, errDeploymentNotFound) + houstonMock.AssertExpectations(s.T()) + s.Equal(returnedDeploymentId, "") + }) + + s.Run("When an error occurs in the GetDeployment api call", func() { + getDeploymentIDForCurrentCommandVar = func(houstonClient houston.ClientInterface, wsID, deploymentID string, prompt bool) (string, []houston.Deployment, error) { + return deploymentID, nil, nil + } + houstonMock := new(houston_mocks.ClientInterface) + houstonMock.On("GetDeployment", mock.Anything).Return(nil, errMockHouston).Once() + + returnedDeploymentId, err := UpdateDeploymentImage(houstonMock, deploymentID, wsID, runtimeVersion, imageName) + s.ErrorContains(err, "failed to get deployment info: some houston error") + houstonMock.AssertExpectations(s.T()) + s.Equal(returnedDeploymentId, "") + }) + + s.Run("Houston API call throws error", func() { + getDeploymentIDForCurrentCommandVar = func(houstonClient houston.ClientInterface, wsID, deploymentID string, prompt bool) (string, []houston.Deployment, error) { + return deploymentID, nil, nil + } + houstonMock := new(houston_mocks.ClientInterface) + deployment := &houston.Deployment{ + ReleaseName: releaseName, + } + houstonMock.On("GetDeployment", mock.Anything).Return(deployment, nil).Once() + houstonMock.On("UpdateDeploymentImage", mock.Anything).Return(nil, errMockHouston).Once() + returnedDeploymentId, err := UpdateDeploymentImage(houstonMock, deploymentID, wsID, runtimeVersion, imageName) + s.ErrorContains(err, "some houston error") + houstonMock.AssertExpectations(s.T()) + s.Equal(returnedDeploymentId, deploymentID) + }) + + s.Run("Successful API call", func() { + getDeploymentIDForCurrentCommandVar = func(houstonClient houston.ClientInterface, wsID, deploymentID string, prompt bool) (string, []houston.Deployment, error) { + return deploymentID, nil, nil + } + houstonMock := new(houston_mocks.ClientInterface) + updateDeploymentImageResp := &houston.UpdateDeploymentImageResp{ + ReleaseName: releaseName, + AirflowVersion: "", + RuntimeVersion: runtimeVersion, + } + deployment := &houston.Deployment{ + ReleaseName: releaseName, + } + houstonMock.On("GetDeployment", mock.Anything).Return(deployment, nil).Once() + houstonMock.On("UpdateDeploymentImage", mock.Anything).Return(updateDeploymentImageResp, nil).Once() + returnedDeploymentId, err := UpdateDeploymentImage(houstonMock, deploymentID, wsID, runtimeVersion, imageName) + houstonMock.AssertExpectations(s.T()) + s.ErrorIs(err, nil) + s.Equal(returnedDeploymentId, deploymentID) + }) +} From d1a9b6a397b49024ec780d4167af07cb9bcda619 Mon Sep 17 00:00:00 2001 From: Rujhan Arora Date: Fri, 13 Dec 2024 17:17:41 +0530 Subject: [PATCH 13/17] Introduction of astro deploy --image-name=image_name --remote --runtime-version=rt_version --- software/deploy/deploy_test.go | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/software/deploy/deploy_test.go b/software/deploy/deploy_test.go index 07dfa328e..5b757939c 100644 --- a/software/deploy/deploy_test.go +++ b/software/deploy/deploy_test.go @@ -1044,10 +1044,10 @@ func (s *Suite) TestUpdateDeploymentImage() { s.Run("When runtimeVersion is empty", func() { houstonMock := new(houston_mocks.ClientInterface) - returnedDeploymentId, err := UpdateDeploymentImage(houstonMock, deploymentID, wsID, "", imageName) + returnedDeploymentID, err := UpdateDeploymentImage(houstonMock, deploymentID, wsID, "", imageName) s.ErrorIs(err, ErrRuntimeVersionNotPassedForRemoteImage) houstonMock.AssertExpectations(s.T()) - s.Equal(returnedDeploymentId, "") + s.Equal(returnedDeploymentID, "") }) s.Run("When getDeploymentIDForCurrentCommandVar gives an error", func() { @@ -1055,10 +1055,10 @@ func (s *Suite) TestUpdateDeploymentImage() { return deploymentID, nil, errDeploymentNotFound } houstonMock := new(houston_mocks.ClientInterface) - returnedDeploymentId, err := UpdateDeploymentImage(houstonMock, deploymentID, wsID, runtimeVersion, imageName) + returnedDeploymentID, err := UpdateDeploymentImage(houstonMock, deploymentID, wsID, runtimeVersion, imageName) s.ErrorIs(err, errDeploymentNotFound) houstonMock.AssertExpectations(s.T()) - s.Equal(returnedDeploymentId, "") + s.Equal(returnedDeploymentID, "") }) s.Run("When an error occurs in the GetDeployment api call", func() { @@ -1068,10 +1068,10 @@ func (s *Suite) TestUpdateDeploymentImage() { houstonMock := new(houston_mocks.ClientInterface) houstonMock.On("GetDeployment", mock.Anything).Return(nil, errMockHouston).Once() - returnedDeploymentId, err := UpdateDeploymentImage(houstonMock, deploymentID, wsID, runtimeVersion, imageName) + returnedDeploymentID, err := UpdateDeploymentImage(houstonMock, deploymentID, wsID, runtimeVersion, imageName) s.ErrorContains(err, "failed to get deployment info: some houston error") houstonMock.AssertExpectations(s.T()) - s.Equal(returnedDeploymentId, "") + s.Equal(returnedDeploymentID, "") }) s.Run("Houston API call throws error", func() { @@ -1084,10 +1084,10 @@ func (s *Suite) TestUpdateDeploymentImage() { } houstonMock.On("GetDeployment", mock.Anything).Return(deployment, nil).Once() houstonMock.On("UpdateDeploymentImage", mock.Anything).Return(nil, errMockHouston).Once() - returnedDeploymentId, err := UpdateDeploymentImage(houstonMock, deploymentID, wsID, runtimeVersion, imageName) + returnedDeploymentID, err := UpdateDeploymentImage(houstonMock, deploymentID, wsID, runtimeVersion, imageName) s.ErrorContains(err, "some houston error") houstonMock.AssertExpectations(s.T()) - s.Equal(returnedDeploymentId, deploymentID) + s.Equal(returnedDeploymentID, deploymentID) }) s.Run("Successful API call", func() { @@ -1105,9 +1105,9 @@ func (s *Suite) TestUpdateDeploymentImage() { } houstonMock.On("GetDeployment", mock.Anything).Return(deployment, nil).Once() houstonMock.On("UpdateDeploymentImage", mock.Anything).Return(updateDeploymentImageResp, nil).Once() - returnedDeploymentId, err := UpdateDeploymentImage(houstonMock, deploymentID, wsID, runtimeVersion, imageName) + returnedDeploymentID, err := UpdateDeploymentImage(houstonMock, deploymentID, wsID, runtimeVersion, imageName) houstonMock.AssertExpectations(s.T()) s.ErrorIs(err, nil) - s.Equal(returnedDeploymentId, deploymentID) + s.Equal(returnedDeploymentID, deploymentID) }) } From cbbf28904f9d6ce58787e3eaf0041ca5820d7bb4 Mon Sep 17 00:00:00 2001 From: Rujhan Arora Date: Mon, 16 Dec 2024 17:56:37 +0530 Subject: [PATCH 14/17] Added setup and teardown functions and moved the code for newly added function there --- cmd/software/deploy.go | 30 +++++++------ cmd/software/deploy_test.go | 25 ++++++++--- cmd/software/suite_test.go | 11 ++++- houston/suite_test.go | 5 +++ software/deploy/deploy.go | 8 ++-- software/deploy/deploy_test.go | 80 +++++++++++++++------------------- 6 files changed, 88 insertions(+), 71 deletions(-) diff --git a/cmd/software/deploy.go b/cmd/software/deploy.go index d4305db89..90bd0134c 100644 --- a/cmd/software/deploy.go +++ b/cmd/software/deploy.go @@ -21,17 +21,18 @@ var ( ignoreCacheDeploy = false - EnsureProjectDir = utils.EnsureProjectDir - DeployAirflowImage = deploy.Airflow - DagsOnlyDeploy = deploy.DagsOnlyDeploy - UpdateDeploymentImage = deploy.UpdateDeploymentImage - isDagOnlyDeploy bool - description string - isImageOnlyDeploy bool - imageName string - runtimeVersionForImageName string - imagePresentOnRemote bool - ErrBothDagsOnlyAndImageOnlySet = errors.New("cannot use both --dags and --image together. Run 'astro deploy' to update both your image and dags") + EnsureProjectDir = utils.EnsureProjectDir + DeployAirflowImage = deploy.Airflow + DagsOnlyDeploy = deploy.DagsOnlyDeploy + UpdateDeploymentImage = deploy.UpdateDeploymentImage + isDagOnlyDeploy bool + description string + isImageOnlyDeploy bool + imageName string + runtimeVersionForImageName string + imagePresentOnRemote bool + ErrBothDagsOnlyAndImageOnlySet = errors.New("cannot use both --dags and --image together. Run 'astro deploy' to update both your image and dags") + ErrImageNameNotPassedForRemoteFlag = errors.New("--image-name is mandatory when --remote flag is passed") ) var deployExample = ` @@ -67,7 +68,7 @@ func NewDeployCmd() *cobra.Command { cmd.Flags().BoolVarP(&isImageOnlyDeploy, "image", "", false, "Push only an image to your Astro Deployment. This only works for Dag-only, Git-sync-based and NFS-based deployments.") cmd.Flags().StringVarP(&imageName, "image-name", "i", "", "Name of the custom image(should be present locally unless --remote is specified) to deploy") cmd.Flags().StringVar(&runtimeVersionForImageName, "runtime-version", "", "Runtime version of the image to deploy. Example - 12.1.1. Mandatory if --image-name --remote is provided") - cmd.Flags().BoolVarP(&imagePresentOnRemote, "remote", "", false, "Custom image is present on remote registry") + cmd.Flags().BoolVarP(&imagePresentOnRemote, "remote", "", false, "Custom image which is present on the remote registry. Can only be used with --image-name flag") if !context.IsCloudContext() && houston.VerifyVersionMatch(houstonVersion, houston.VersionRestrictions{GTE: "0.34.0"}) { cmd.Flags().BoolVarP(&isDagOnlyDeploy, "dags", "d", false, "Push only DAGs to your Deployment") @@ -126,7 +127,10 @@ func deployAirflow(cmd *cobra.Command, args []string) error { return DagsOnlyDeploy(houstonClient, appConfig, ws, deploymentID, config.WorkingPath, nil, true, description) } - if imagePresentOnRemote && imageName != "" { + if imagePresentOnRemote { + if imageName == "" { + return ErrImageNameNotPassedForRemoteFlag + } deploymentID, err = UpdateDeploymentImage(houstonClient, deploymentID, ws, runtimeVersionForImageName, imageName) if err != nil { return err diff --git a/cmd/software/deploy_test.go b/cmd/software/deploy_test.go index 0ad9ba426..7775cdb24 100644 --- a/cmd/software/deploy_test.go +++ b/cmd/software/deploy_test.go @@ -63,9 +63,6 @@ func (s *Suite) TestDeploy() { err = execDeployCmd([]string{"test-deployment-id", "--force"}...) s.NoError(err) - // Restore DagsOnlyDeploy to default behavior - DagsOnlyDeploy = deploy.DagsOnlyDeploy - s.Run("error should be returned for astro deploy, if DeployAirflowImage throws error", func() { DeployAirflowImage = func(houstonClient houston.ClientInterface, path, deploymentID, wsID, byoRegistryDomain string, ignoreCacheDeploy, byoRegistryEnabled, prompt bool, description string, isImageOnlyDeploy bool, imageName string) (string, error) { return deploymentID, deploy.ErrNoWorkspaceID @@ -85,7 +82,6 @@ func (s *Suite) TestDeploy() { } err := execDeployCmd([]string{"-f"}...) s.ErrorIs(err, deploy.ErrNoWorkspaceID) - DagsOnlyDeploy = deploy.DagsOnlyDeploy }) s.Run("No error should be returned for astro deploy, if dags deploy throws error but the feature itself is disabled", func() { @@ -143,13 +139,22 @@ func (s *Suite) TestDeploy() { DagsOnlyDeploy = func(houstonClient houston.ClientInterface, appConfig *houston.AppConfig, wsID, deploymentID, dagsParentPath string, dagDeployURL *string, cleanUpFiles bool, description string) error { return nil } + // Create a flag to track if DeployAirflowImage is called + deployAirflowImageCalled := false + + // Mock function for DeployAirflowImage + DeployAirflowImage = func(houstonClient houston.ClientInterface, path, deploymentID, wsID, byoRegistryDomain string, ignoreCacheDeploy, byoRegistryEnabled, prompt bool, description string, isImageOnlyDeploy bool, imageName string) (string, error) { + deployAirflowImageCalled = true // Set the flag if this function is called + return deploymentID, nil + } UpdateDeploymentImage = func(houstonClient houston.ClientInterface, deploymentID, wsID, runtimeVersion, imageName string) (string, error) { return "", nil } testImageName := "test-image-name" // Set the expected image name err := execDeployCmd([]string{"test-deployment-id", "--image-name=" + testImageName, "--force", "--remote", "--workspace-id=" + mockWorkspace.ID}...) s.ErrorIs(err, nil) - UpdateDeploymentImage = deploy.UpdateDeploymentImage + // Assert that DeployAirflowImage was NOT called + s.False(deployAirflowImageCalled, "DeployAirflowImage should not be called when --remote is specified") }) s.Run("Test for the flag --image-name with --remote. Dags should not be deployed if UpdateDeploymentImage throws an error", func() { @@ -159,7 +164,14 @@ func (s *Suite) TestDeploy() { testImageName := "test-image-name" // Set the expected image name err := execDeployCmd([]string{"test-deployment-id", "--image-name=" + testImageName, "--force", "--remote", "--workspace-id=" + mockWorkspace.ID}...) s.ErrorIs(err, errNoWorkspaceFound) - UpdateDeploymentImage = deploy.UpdateDeploymentImage + }) + + s.Run("Test for the flag --remote without --image-name. It should throw an error", func() { + UpdateDeploymentImage = func(houstonClient houston.ClientInterface, deploymentID, wsID, runtimeVersion, imageName string) (string, error) { + return "", errNoWorkspaceFound + } + err := execDeployCmd([]string{"test-deployment-id", "--force", "--remote", "--workspace-id=" + mockWorkspace.ID}...) + s.ErrorIs(err, ErrImageNameNotPassedForRemoteFlag) }) s.Run("error should be returned if BYORegistryEnabled is true but BYORegistryDomain is empty", func() { @@ -174,6 +186,5 @@ func (s *Suite) TestDeploy() { } err := execDeployCmd([]string{"-f"}...) s.ErrorIs(err, deploy.ErrBYORegistryDomainNotSet) - DagsOnlyDeploy = deploy.DagsOnlyDeploy }) } diff --git a/cmd/software/suite_test.go b/cmd/software/suite_test.go index 0071f2e6a..20c29742e 100644 --- a/cmd/software/suite_test.go +++ b/cmd/software/suite_test.go @@ -4,6 +4,7 @@ import ( "testing" testUtil "github.com/astronomer/astro-cli/pkg/testing" + "github.com/astronomer/astro-cli/software/deploy" "github.com/stretchr/testify/suite" ) @@ -24,7 +25,13 @@ func (s *Suite) SetupTest() { houstonVersion = "0.34.0" } +func (s *Suite) TearDownSuite() { + DagsOnlyDeploy = deploy.DagsOnlyDeploy + UpdateDeploymentImage = deploy.UpdateDeploymentImage +} + var ( - _ suite.SetupAllSuite = (*Suite)(nil) - _ suite.SetupTestSuite = (*Suite)(nil) + _ suite.SetupAllSuite = (*Suite)(nil) + _ suite.SetupTestSuite = (*Suite)(nil) + _ suite.TearDownAllSuite = (*Suite)(nil) ) diff --git a/houston/suite_test.go b/houston/suite_test.go index 45293345d..26af464b8 100644 --- a/houston/suite_test.go +++ b/houston/suite_test.go @@ -3,6 +3,7 @@ package houston import ( "testing" + testUtil "github.com/astronomer/astro-cli/pkg/testing" "github.com/stretchr/testify/suite" ) @@ -10,6 +11,10 @@ type Suite struct { suite.Suite } +func SetupSuite() { + testUtil.InitTestConfig(testUtil.SoftwarePlatform) +} + func TestHouston(t *testing.T) { suite.Run(t, new(Suite)) } diff --git a/software/deploy/deploy.go b/software/deploy/deploy.go index 2eb8a9251..babeff618 100644 --- a/software/deploy/deploy.go +++ b/software/deploy/deploy.go @@ -172,11 +172,10 @@ func UpdateDeploymentImage(houstonClient houston.ClientInterface, deploymentID, if runtimeVersion == "" { return "", ErrRuntimeVersionNotPassedForRemoteImage } - deploymentIDForCurrentCmd, _, err := getDeploymentIDForCurrentCommandVar(houstonClient, wsID, deploymentID, deploymentID == "") + deploymentID, _, err := getDeploymentIDForCurrentCommandVar(houstonClient, wsID, deploymentID, deploymentID == "") if err != nil { return "", err } - deploymentID = deploymentIDForCurrentCmd if deploymentID == "" { return "", errInvalidDeploymentID } @@ -187,7 +186,7 @@ func UpdateDeploymentImage(houstonClient houston.ClientInterface, deploymentID, fmt.Println("Skipping building the image since --image-name flag is used...") _, err = updateDeploymentImageAPICall(houstonClient, imageName, deploymentInfo.ReleaseName, "", runtimeVersion) fmt.Println("Image successfully updated") - return deploymentIDForCurrentCmd, err + return deploymentID, err } func pushDockerImage(byoRegistryEnabled bool, byoRegistryDomain, name, nextTag, cloudDomain string, imageHandler airflow.ImageHandler, houstonClient houston.ClientInterface, c *config.Context) error { @@ -429,11 +428,10 @@ func DagsOnlyDeploy(houstonClient houston.ClientInterface, appConfig *houston.Ap return ErrDagOnlyDeployDisabledInConfig } - deploymentIDForCurrentCmd, _, err := getDeploymentIDForCurrentCommandVar(houstonClient, wsID, deploymentID, deploymentID == "") + deploymentID, _, err := getDeploymentIDForCurrentCommandVar(houstonClient, wsID, deploymentID, deploymentID == "") if err != nil { return err } - deploymentID = deploymentIDForCurrentCmd if deploymentID == "" { return errInvalidDeploymentID diff --git a/software/deploy/deploy_test.go b/software/deploy/deploy_test.go index 5b757939c..a25970320 100644 --- a/software/deploy/deploy_test.go +++ b/software/deploy/deploy_test.go @@ -113,6 +113,7 @@ func (s *Suite) SetupSuite() { } func (s *Suite) SetupTest() { + testUtil.InitTestConfig(testUtil.SoftwarePlatform) s.mockImageHandler = new(mocks.ImageHandler) imageHandlerInit = func(image string) airflow.ImageHandler { return s.mockImageHandler @@ -120,8 +121,21 @@ func (s *Suite) SetupTest() { s.houstonMock = new(houston_mocks.ClientInterface) } +func (s *Suite) SetupSubTest() { + testUtil.InitTestConfig(testUtil.SoftwarePlatform) + s.mockImageHandler = new(mocks.ImageHandler) + imageHandlerInit = func(image string) airflow.ImageHandler { + return s.mockImageHandler + } + s.houstonMock = new(houston_mocks.ClientInterface) +} + +func (s *Suite) TearDownSubTest() { + s.houstonMock.AssertExpectations(s.T()) + s.mockImageHandler.AssertExpectations(s.T()) +} + func (s *Suite) TearDownSuite() { - // Cleanup logic, if any (e.g., clearing mocks) s.mockImageHandler = nil s.houstonMock = nil s.fsForDockerConfig = nil @@ -290,10 +304,7 @@ func (s *Suite) TestBuildPushDockerImageFailure() { s.EqualError(err, "failed to parse dockerfile: testfiles/Dockerfile.invalid: when using JSON array syntax, arrays must be comprised of strings only") dockerfile = "Dockerfile" - fs := afero.NewMemMapFs() - configYaml := testUtil.NewTestConfig("docker") - afero.WriteFile(fs, config.HomeConfigFile, configYaml, 0o777) - config.InitConfig(fs) + config.InitConfig(s.fsForDockerConfig) mockedDeploymentConfig := &houston.DeploymentConfig{ AirflowImages: mockAirflowImageList, @@ -459,16 +470,12 @@ func (s *Suite) TestAirflowSuccess() { } func (s *Suite) TestAirflowSuccessForImageOnly() { - fs := afero.NewMemMapFs() - configYaml := testUtil.NewTestConfig("localhost") - afero.WriteFile(fs, config.HomeConfigFile, configYaml, 0o777) - config.InitConfig(fs) + config.InitConfig(s.fsForLocalConfig) - mockImageHandler := new(mocks.ImageHandler) imageHandlerInit = func(image string) airflow.ImageHandler { - mockImageHandler.On("Build", mock.Anything, mock.Anything, mock.Anything).Return(nil) - mockImageHandler.On("Push", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil) - return mockImageHandler + s.mockImageHandler.On("Build", mock.Anything, mock.Anything, mock.Anything).Return(nil) + s.mockImageHandler.On("Push", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil) + return s.mockImageHandler } mockedDeploymentConfig := &houston.DeploymentConfig{ @@ -478,10 +485,9 @@ func (s *Suite) TestAirflowSuccessForImageOnly() { houston.RuntimeRelease{Version: "4.2.4", AirflowVersion: "2.2.5"}, houston.RuntimeRelease{Version: "4.2.5", AirflowVersion: "2.2.5"}, } - houstonMock := new(houston_mocks.ClientInterface) - houstonMock.On("GetWorkspace", mock.Anything).Return(&houston.Workspace{}, nil).Once() - houstonMock.On("ListDeployments", mock.Anything).Return([]houston.Deployment{{ID: "test-deployment-id"}}, nil).Once() - houstonMock.On("GetDeploymentConfig", nil).Return(mockedDeploymentConfig, nil).Once() + s.houstonMock.On("GetWorkspace", mock.Anything).Return(&houston.Workspace{}, nil).Once() + s.houstonMock.On("ListDeployments", mock.Anything).Return([]houston.Deployment{{ID: "test-deployment-id"}}, nil).Once() + s.houstonMock.On("GetDeploymentConfig", nil).Return(mockedDeploymentConfig, nil).Once() dagDeployment := &houston.DagDeploymentConfig{ Type: "dag-only", } @@ -489,13 +495,11 @@ func (s *Suite) TestAirflowSuccessForImageOnly() { DagDeployment: *dagDeployment, } - houstonMock.On("GetDeployment", mock.Anything).Return(deployment, nil).Once() - houstonMock.On("GetRuntimeReleases", "").Return(mockRuntimeReleases, nil) + s.houstonMock.On("GetDeployment", mock.Anything).Return(deployment, nil).Once() + s.houstonMock.On("GetRuntimeReleases", "").Return(mockRuntimeReleases, nil) - _, err := Airflow(houstonMock, "./testfiles/", "test-deployment-id", "test-workspace-id", "", false, false, false, description, true, "") + _, err := Airflow(s.houstonMock, "./testfiles/", "test-deployment-id", "test-workspace-id", "", false, false, false, description, true, "") s.NoError(err) - houstonMock.AssertExpectations(s.T()) - mockImageHandler.AssertExpectations(s.T()) } func (s *Suite) TestAirflowSuccessForImageName() { @@ -540,7 +544,6 @@ func (s *Suite) TestAirflowFailForImageNameWhenImageHasNoRuntimeLabel() { s.mockImageHandler.On("GetLabel", "", airflow.RuntimeImageLabel).Return("", nil) return s.mockImageHandler } - s.houstonMock.On("GetWorkspace", mock.Anything).Return(&houston.Workspace{}, nil).Once() s.houstonMock.On("ListDeployments", mock.Anything).Return([]houston.Deployment{{ID: "test-deployment-id"}}, nil).Once() dagDeployment := &houston.DagDeploymentConfig{ @@ -1035,7 +1038,6 @@ func (s *Suite) TestDeployDagsOnlyFailure() { } func (s *Suite) TestUpdateDeploymentImage() { - testUtil.InitTestConfig(testUtil.SoftwarePlatform) deploymentID := "test-deployment-id" wsID := "test-workspace-id" runtimeVersion := "12.1.1" @@ -1043,10 +1045,8 @@ func (s *Suite) TestUpdateDeploymentImage() { releaseName := "releaseName" s.Run("When runtimeVersion is empty", func() { - houstonMock := new(houston_mocks.ClientInterface) - returnedDeploymentID, err := UpdateDeploymentImage(houstonMock, deploymentID, wsID, "", imageName) + returnedDeploymentID, err := UpdateDeploymentImage(s.houstonMock, deploymentID, wsID, "", imageName) s.ErrorIs(err, ErrRuntimeVersionNotPassedForRemoteImage) - houstonMock.AssertExpectations(s.T()) s.Equal(returnedDeploymentID, "") }) @@ -1054,10 +1054,8 @@ func (s *Suite) TestUpdateDeploymentImage() { getDeploymentIDForCurrentCommandVar = func(houstonClient houston.ClientInterface, wsID, deploymentID string, prompt bool) (string, []houston.Deployment, error) { return deploymentID, nil, errDeploymentNotFound } - houstonMock := new(houston_mocks.ClientInterface) - returnedDeploymentID, err := UpdateDeploymentImage(houstonMock, deploymentID, wsID, runtimeVersion, imageName) + returnedDeploymentID, err := UpdateDeploymentImage(s.houstonMock, deploymentID, wsID, runtimeVersion, imageName) s.ErrorIs(err, errDeploymentNotFound) - houstonMock.AssertExpectations(s.T()) s.Equal(returnedDeploymentID, "") }) @@ -1065,12 +1063,10 @@ func (s *Suite) TestUpdateDeploymentImage() { getDeploymentIDForCurrentCommandVar = func(houstonClient houston.ClientInterface, wsID, deploymentID string, prompt bool) (string, []houston.Deployment, error) { return deploymentID, nil, nil } - houstonMock := new(houston_mocks.ClientInterface) - houstonMock.On("GetDeployment", mock.Anything).Return(nil, errMockHouston).Once() + s.houstonMock.On("GetDeployment", mock.Anything).Return(nil, errMockHouston).Once() - returnedDeploymentID, err := UpdateDeploymentImage(houstonMock, deploymentID, wsID, runtimeVersion, imageName) + returnedDeploymentID, err := UpdateDeploymentImage(s.houstonMock, deploymentID, wsID, runtimeVersion, imageName) s.ErrorContains(err, "failed to get deployment info: some houston error") - houstonMock.AssertExpectations(s.T()) s.Equal(returnedDeploymentID, "") }) @@ -1078,15 +1074,13 @@ func (s *Suite) TestUpdateDeploymentImage() { getDeploymentIDForCurrentCommandVar = func(houstonClient houston.ClientInterface, wsID, deploymentID string, prompt bool) (string, []houston.Deployment, error) { return deploymentID, nil, nil } - houstonMock := new(houston_mocks.ClientInterface) deployment := &houston.Deployment{ ReleaseName: releaseName, } - houstonMock.On("GetDeployment", mock.Anything).Return(deployment, nil).Once() - houstonMock.On("UpdateDeploymentImage", mock.Anything).Return(nil, errMockHouston).Once() - returnedDeploymentID, err := UpdateDeploymentImage(houstonMock, deploymentID, wsID, runtimeVersion, imageName) + s.houstonMock.On("GetDeployment", mock.Anything).Return(deployment, nil).Once() + s.houstonMock.On("UpdateDeploymentImage", mock.Anything).Return(nil, errMockHouston).Once() + returnedDeploymentID, err := UpdateDeploymentImage(s.houstonMock, deploymentID, wsID, runtimeVersion, imageName) s.ErrorContains(err, "some houston error") - houstonMock.AssertExpectations(s.T()) s.Equal(returnedDeploymentID, deploymentID) }) @@ -1094,7 +1088,6 @@ func (s *Suite) TestUpdateDeploymentImage() { getDeploymentIDForCurrentCommandVar = func(houstonClient houston.ClientInterface, wsID, deploymentID string, prompt bool) (string, []houston.Deployment, error) { return deploymentID, nil, nil } - houstonMock := new(houston_mocks.ClientInterface) updateDeploymentImageResp := &houston.UpdateDeploymentImageResp{ ReleaseName: releaseName, AirflowVersion: "", @@ -1103,10 +1096,9 @@ func (s *Suite) TestUpdateDeploymentImage() { deployment := &houston.Deployment{ ReleaseName: releaseName, } - houstonMock.On("GetDeployment", mock.Anything).Return(deployment, nil).Once() - houstonMock.On("UpdateDeploymentImage", mock.Anything).Return(updateDeploymentImageResp, nil).Once() - returnedDeploymentID, err := UpdateDeploymentImage(houstonMock, deploymentID, wsID, runtimeVersion, imageName) - houstonMock.AssertExpectations(s.T()) + s.houstonMock.On("GetDeployment", mock.Anything).Return(deployment, nil).Once() + s.houstonMock.On("UpdateDeploymentImage", mock.Anything).Return(updateDeploymentImageResp, nil).Once() + returnedDeploymentID, err := UpdateDeploymentImage(s.houstonMock, deploymentID, wsID, runtimeVersion, imageName) s.ErrorIs(err, nil) s.Equal(returnedDeploymentID, deploymentID) }) From dc528341e83013308da8b303542592241b4904f0 Mon Sep 17 00:00:00 2001 From: Rujhan Arora Date: Mon, 16 Dec 2024 18:09:28 +0530 Subject: [PATCH 15/17] Fixed tests --- cmd/software/deploy_test.go | 3 +++ cmd/software/suite_test.go | 5 +++++ 2 files changed, 8 insertions(+) diff --git a/cmd/software/deploy_test.go b/cmd/software/deploy_test.go index 7775cdb24..0b21fb15a 100644 --- a/cmd/software/deploy_test.go +++ b/cmd/software/deploy_test.go @@ -63,6 +63,9 @@ func (s *Suite) TestDeploy() { err = execDeployCmd([]string{"test-deployment-id", "--force"}...) s.NoError(err) + // Restore DagsOnlyDeploy to default behavior + DagsOnlyDeploy = deploy.DagsOnlyDeploy + s.Run("error should be returned for astro deploy, if DeployAirflowImage throws error", func() { DeployAirflowImage = func(houstonClient houston.ClientInterface, path, deploymentID, wsID, byoRegistryDomain string, ignoreCacheDeploy, byoRegistryEnabled, prompt bool, description string, isImageOnlyDeploy bool, imageName string) (string, error) { return deploymentID, deploy.ErrNoWorkspaceID diff --git a/cmd/software/suite_test.go b/cmd/software/suite_test.go index 20c29742e..d7fe4cac1 100644 --- a/cmd/software/suite_test.go +++ b/cmd/software/suite_test.go @@ -23,6 +23,11 @@ func (s *Suite) SetupSuite() { func (s *Suite) SetupTest() { // Reset the version once this is torn down houstonVersion = "0.34.0" + DagsOnlyDeploy = deploy.DagsOnlyDeploy +} + +func (s *Suite) SetupSubTest() { + DagsOnlyDeploy = deploy.DagsOnlyDeploy } func (s *Suite) TearDownSuite() { From 37939cc60948f5bf5b2417729758804c4f8683f5 Mon Sep 17 00:00:00 2001 From: Rujhan Arora Date: Tue, 17 Dec 2024 14:04:03 +0530 Subject: [PATCH 16/17] Removed the abstracted out updateDeploymentImageAPICall function --- software/deploy/deploy.go | 11 ++++------- 1 file changed, 4 insertions(+), 7 deletions(-) diff --git a/software/deploy/deploy.go b/software/deploy/deploy.go index babeff618..9912b24f5 100644 --- a/software/deploy/deploy.go +++ b/software/deploy/deploy.go @@ -163,11 +163,6 @@ func validateRuntimeVersion(houstonClient houston.ClientInterface, tag string, d return nil } -func updateDeploymentImageAPICall(houstonClient houston.ClientInterface, imageName, releaseName, airflowVersion, runtimeVersion string) (interface{}, error) { - req := houston.UpdateDeploymentImageRequest{ReleaseName: releaseName, Image: imageName, AirflowVersion: airflowVersion, RuntimeVersion: runtimeVersion} - return houston.Call(houstonClient.UpdateDeploymentImage)(req) -} - func UpdateDeploymentImage(houstonClient houston.ClientInterface, deploymentID, wsID, runtimeVersion, imageName string) (string, error) { if runtimeVersion == "" { return "", ErrRuntimeVersionNotPassedForRemoteImage @@ -184,7 +179,8 @@ func UpdateDeploymentImage(houstonClient houston.ClientInterface, deploymentID, return "", fmt.Errorf("failed to get deployment info: %w", err) } fmt.Println("Skipping building the image since --image-name flag is used...") - _, err = updateDeploymentImageAPICall(houstonClient, imageName, deploymentInfo.ReleaseName, "", runtimeVersion) + req := houston.UpdateDeploymentImageRequest{ReleaseName: deploymentInfo.ReleaseName, Image: imageName, AirflowVersion: "", RuntimeVersion: runtimeVersion} + _, err = houston.Call(houstonClient.UpdateDeploymentImage)(req) fmt.Println("Image successfully updated") return deploymentID, err } @@ -214,7 +210,8 @@ func pushDockerImage(byoRegistryEnabled bool, byoRegistryDomain, name, nextTag, } runtimeVersion, _ := imageHandler.GetLabel("", runtimeImageLabel) airflowVersion, _ := imageHandler.GetLabel("", airflowImageLabel) - _, err := updateDeploymentImageAPICall(houstonClient, remoteImage, name, airflowVersion, runtimeVersion) + req := houston.UpdateDeploymentImageRequest{ReleaseName: name, Image: remoteImage, AirflowVersion: airflowVersion, RuntimeVersion: runtimeVersion} + _, err = houston.Call(houstonClient.UpdateDeploymentImage)(req) return err } return nil From 21dec01b4e83de70e9bfacc0591365c9a15bcf0e Mon Sep 17 00:00:00 2001 From: Rujhan Arora Date: Tue, 17 Dec 2024 14:15:48 +0530 Subject: [PATCH 17/17] Removed already added TearDownSubTest --- software/deploy/deploy_test.go | 5 ----- 1 file changed, 5 deletions(-) diff --git a/software/deploy/deploy_test.go b/software/deploy/deploy_test.go index a25970320..acefe727b 100644 --- a/software/deploy/deploy_test.go +++ b/software/deploy/deploy_test.go @@ -143,11 +143,6 @@ func (s *Suite) TearDownSuite() { imageHandlerInit = airflow.ImageHandlerInit } -func (s *Suite) TearDownSubTest() { - s.houstonMock.AssertExpectations(s.T()) - s.mockImageHandler.AssertExpectations(s.T()) -} - func (s *Suite) TearDownTest() { s.houstonMock.AssertExpectations(s.T()) s.mockImageHandler.AssertExpectations(s.T())