Skip to content

Commit

Permalink
Added setup and teardown functions and moved the code for newly added…
Browse files Browse the repository at this point in the history 
… function there
  • Loading branch information
@rujhan-arora-astronomer
rujhan-arora-astronomer committed Dec 16, 2024
1 parent cc9dae8 commit e55b75c
Show file tree
Hide file tree
Showing 6 changed files with 88 additions and 71 deletions.
30 changes: 17 additions & 13 deletions cmd/software/deploy.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 = `
Expand Down Expand Up @@ -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")
Expand Down Expand Up @@ -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
Expand Down
25 changes: 18 additions & 7 deletions cmd/software/deploy_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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() {
Expand Down Expand Up @@ -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() {
Expand All @@ -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() {
Expand All @@ -174,6 +186,5 @@ func (s *Suite) TestDeploy() {
}
err := execDeployCmd([]string{"-f"}...)
s.ErrorIs(err, deploy.ErrBYORegistryDomainNotSet)
DagsOnlyDeploy = deploy.DagsOnlyDeploy
})
}
11 changes: 9 additions & 2 deletions cmd/software/suite_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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"
)

Expand All @@ -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)
)
5 changes: 5 additions & 0 deletions houston/suite_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,13 +3,18 @@ package houston
import (
"testing"

testUtil "github.com/astronomer/astro-cli/pkg/testing"
"github.com/stretchr/testify/suite"
)

type Suite struct {
suite.Suite
}

func SetupSuite() {
testUtil.InitTestConfig(testUtil.SoftwarePlatform)
}

func TestHouston(t *testing.T) {
suite.Run(t, new(Suite))
}
8 changes: 3 additions & 5 deletions software/deploy/deploy.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
Expand All @@ -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 {
Expand Down Expand Up @@ -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
Expand Down
80 changes: 36 additions & 44 deletions software/deploy/deploy_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -113,15 +113,29 @@ 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
}
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
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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{
Expand All @@ -468,24 +475,21 @@ 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",
}
deployment := &houston.Deployment{
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() {
Expand Down Expand Up @@ -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{
Expand Down Expand Up @@ -1031,66 +1034,56 @@ func (s *Suite) TestDeployDagsOnlyFailure() {
}

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)
returnedDeploymentID, err := UpdateDeploymentImage(s.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)
returnedDeploymentID, err := UpdateDeploymentImage(s.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()
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, "")
})

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.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)
})

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: "",
Expand All @@ -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)
})
Expand Down

0 comments on commit e55b75c

Please sign in to comment.