From e55b75c9366d4a68f13b14db38c4a48c0d633f55 Mon Sep 17 00:00:00 2001 From: Rujhan Arora Date: Mon, 16 Dec 2024 17:56:37 +0530 Subject: [PATCH] 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 09bdc29b5..a1f2bad1a 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 7349c47c7..5c6c42a31 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 @@ -280,10 +294,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, @@ -449,16 +460,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{ @@ -468,10 +475,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", } @@ -479,13 +485,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() { @@ -532,7 +536,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{ @@ -1031,7 +1034,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" @@ -1039,10 +1041,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, "") }) @@ -1050,10 +1050,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, "") }) @@ -1061,12 +1059,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, "") }) @@ -1074,15 +1070,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) }) @@ -1090,7 +1084,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: "", @@ -1099,10 +1092,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) })