From ff318060a30e81892c9d6f83f6d494dfbefe38a1 Mon Sep 17 00:00:00 2001 From: David Koenitzer Date: Tue, 12 Sep 2023 09:28:17 -0400 Subject: [PATCH] Fix Upgrade Test (#1385) * fix runtime * update parse test * update mocks * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix lint * fix test * add tests * add tests * fix lint * add test * fix test * fix test * add comment --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .gitignore | 1 + airflow-client/mocks/Client.go | 11 +- airflow/container.go | 1 + airflow/docker.go | 54 ++++++++-- airflow/docker_image.go | 12 +++ airflow/docker_image_test.go | 30 ++++++ airflow/docker_test.go | 117 +++++++++++++++++++++ airflow/include/dagintegritytestdefault.py | 28 +++-- airflow/mocks/ContainerHandler.go | 11 +- airflow/mocks/DockerCLIClient.go | 11 +- airflow/mocks/DockerComposeAPI.go | 11 +- airflow/mocks/DockerRegistryAPI.go | 11 +- airflow/mocks/ImageHandler.go | 25 +++-- airflow/mocks/RegistryHandler.go | 11 +- astro-client-core/mocks/client.go | 11 +- astro-client/mocks/Client.go | 11 +- cmd/airflow.go | 43 +++++--- cmd/airflow_test.go | 60 +++++++++++ cmd/errors.go | 6 +- houston/mocks/ClientInterface.go | 11 +- pkg/azure/mocks/Azure.go | 11 +- 21 files changed, 395 insertions(+), 92 deletions(-) diff --git a/.gitignore b/.gitignore index 5e9519622..086a8a6fc 100644 --- a/.gitignore +++ b/.gitignore @@ -21,4 +21,5 @@ requirements.txt sagemaker-batch-inference.py testbin/ testbin/golangci-lint +upgrade-test* vendor/github.com/theupdateframework/notary/Dockerfile diff --git a/airflow-client/mocks/Client.go b/airflow-client/mocks/Client.go index d60aa5a63..f93ba9e76 100644 --- a/airflow-client/mocks/Client.go +++ b/airflow-client/mocks/Client.go @@ -1,4 +1,4 @@ -// Code generated by mockery v2.20.2. DO NOT EDIT. +// Code generated by mockery v2.32.0. DO NOT EDIT. package airflow_mocks @@ -168,13 +168,12 @@ func (_m *Client) UpdateVariable(airflowURL string, variable airflowclient.Varia return r0 } -type mockConstructorTestingTNewClient interface { +// NewClient creates a new instance of Client. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations. +// The first argument is typically a *testing.T value. +func NewClient(t interface { mock.TestingT Cleanup(func()) -} - -// NewClient creates a new instance of Client. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations. -func NewClient(t mockConstructorTestingTNewClient) *Client { +}) *Client { mock := &Client{} mock.Mock.Test(t) diff --git a/airflow/container.go b/airflow/container.go index c9dbd73e0..518a8701d 100644 --- a/airflow/container.go +++ b/airflow/container.go @@ -48,6 +48,7 @@ type ImageHandler interface { Push(registry, username, token, remoteImage string) error Pull(registry, username, token, remoteImage string) error GetLabel(altImageName, labelName string) (string, error) + DoesImageExist(image string) error ListLabels() (map[string]string, error) TagLocalImage(localImage string) error Run(dagID, envFile, settingsFile, containerName, dagFile, executionDate string, taskLogs bool) error diff --git a/airflow/docker.go b/airflow/docker.go index 7f5516628..d0ec63aa8 100644 --- a/airflow/docker.go +++ b/airflow/docker.go @@ -79,9 +79,10 @@ const ( ) var ( - errNoFile = errors.New("file specified does not exist") - errSettingsPath = "error looking for settings.yaml" - errComposeProjectRunning = errors.New("project is up and running") + errNoFile = errors.New("file specified does not exist") + errSettingsPath = "error looking for settings.yaml" + errComposeProjectRunning = errors.New("project is up and running") + errCustomImageDoesNotExist = errors.New("The custom image provided either does not exist or Docker is unable to connect to the repository") initSettings = settings.ConfigSettings exportSettings = settings.Export @@ -540,8 +541,16 @@ func (d *DockerCompose) UpgradeTest(newAirflowVersion, deploymentID, newImageNam versionTest = true dagTest = true } - // if user supplies deployment id pull down current image var deploymentImage string + // if custom image is used get new Airflow version + if customImage != "" { + err := d.imageHandler.DoesImageExist(customImage) + if err != nil { + return errCustomImageDoesNotExist + } + newAirflowVersion = strings.SplitN(customImage, ":", partsNum)[1] + } + // if user supplies deployment id pull down current image if deploymentID != "" { err := d.pullImageFromDeployment(deploymentID, client) if err != nil { @@ -747,7 +756,7 @@ func (d *DockerCompose) dagTest(testHomeDirectory, newAirflowVersion, newDockerF // create html report htmlReportArgs := "--html=dag-test-report.html --self-contained-html" // compare pip freeze files - fmt.Println("\nRunning parse test") + fmt.Println("\nRunning DAG parse test with the new Airflow version") exitCode, err := d.imageHandler.Pytest(pytestFile, d.airflowHome, d.envFile, testHomeDirectory, strings.Fields(htmlReportArgs), true, airflowTypes.ImageBuildConfig{Path: d.airflowHome, Output: true}) if err != nil { if strings.Contains(exitCode, "1") { // exit code is 1 meaning tests failed @@ -771,7 +780,7 @@ func GetRegistryURL(domain string) string { return registry } -func upgradeDockerfile(oldDockerfilePath, newDockerfilePath, newTag, newImage string) error { +func upgradeDockerfile(oldDockerfilePath, newDockerfilePath, newTag, newImage string) error { //nolint:gocognit // Read the content of the old Dockerfile content, err := os.ReadFile(oldDockerfilePath) if err != nil { @@ -789,6 +798,25 @@ func upgradeDockerfile(oldDockerfilePath, newDockerfilePath, newTag, newImage st line = parts[0] + ":" + newTag } } + if strings.HasPrefix(strings.TrimSpace(line), "FROM quay.io/astronomer/ap-airflow:") { + isRuntime, err := isRuntimeVersion(newTag) + if err != nil { + logrus.Debug(err) + } + if isRuntime { + // Replace the tag on the matching line + parts := strings.SplitN(line, "/", partsNum) + if len(parts) >= partsNum { + line = parts[0] + "/astronomer/astro-runtime:" + newTag + } + } else { + // Replace the tag on the matching line + parts := strings.SplitN(line, ":", partsNum) + if len(parts) == partsNum { + line = parts[0] + ":" + newTag + } + } + } newContent.WriteString(line + "\n") // Add a newline after each line } } else { @@ -813,6 +841,20 @@ func upgradeDockerfile(oldDockerfilePath, newDockerfilePath, newTag, newImage st return nil } +func isRuntimeVersion(versionStr string) (bool, error) { + // Parse the version string + v, err := semver.NewVersion(versionStr) + if err != nil { + return false, err + } + + // Runtime versions start 4.0.0 to not get confused with Airflow versions that are currently in 2.X.X + referenceVersion := semver.MustParse("4.0.0") + + // Compare the parsed version with the reference version + return v.Compare(referenceVersion) > 0, nil +} + func CreateVersionTestFile(beforeFile, afterFile, outputFile string) error { // Open the before file for reading before, err := os.Open(beforeFile) diff --git a/airflow/docker_image.go b/airflow/docker_image.go index 1b54e1585..6d1a766d2 100644 --- a/airflow/docker_image.go +++ b/airflow/docker_image.go @@ -407,6 +407,18 @@ func (d *DockerImage) GetLabel(altImageName, labelName string) (string, error) { return label, nil } +func (d *DockerImage) DoesImageExist(image string) error { + dockerCommand := config.CFG.DockerCommand.GetString() + stdout := new(bytes.Buffer) + stderr := new(bytes.Buffer) + + err := cmdExec(dockerCommand, stdout, stderr, "manifest", "inspect", image) + if err != nil { + return err + } + return nil +} + func (d *DockerImage) ListLabels() (map[string]string, error) { dockerCommand := config.CFG.DockerCommand.GetString() diff --git a/airflow/docker_image_test.go b/airflow/docker_image_test.go index 9f850519b..40d7199fe 100644 --- a/airflow/docker_image_test.go +++ b/airflow/docker_image_test.go @@ -501,6 +501,36 @@ func TestDockerImageListLabel(t *testing.T) { }) } +func TestDoesImageExist(t *testing.T) { + handler := DockerImage{ + imageName: "testing", + } + testImage := "image" + + previousCmdExec := cmdExec + defer func() { cmdExec = previousCmdExec }() + + t.Run("success", func(t *testing.T) { + cmdExec = func(cmd string, stdout, stderr io.Writer, args ...string) error { + assert.Contains(t, args, "inspect") + return nil + } + + err := handler.DoesImageExist(testImage) + assert.NoError(t, err) + }) + + t.Run("cmdExec error", func(t *testing.T) { + cmdExec = func(cmd string, stdout, stderr io.Writer, args ...string) error { + assert.Contains(t, args, "inspect") + return errMockDocker + } + + err := handler.DoesImageExist(testImage) + assert.ErrorIs(t, err, errMockDocker) + }) +} + func TestDockerTagLocalImage(t *testing.T) { handler := DockerImage{ imageName: "testing", diff --git a/airflow/docker_test.go b/airflow/docker_test.go index cac4c22d2..8503e9748 100644 --- a/airflow/docker_test.go +++ b/airflow/docker_test.go @@ -1979,3 +1979,120 @@ func TestCreateDockerProject(t *testing.T) { assert.Equal(t, 5433, int(postgresService.Ports[len(prj.Services[0].Ports)-1].Published)) }) } + +func TestUpgradeDockerfile(t *testing.T) { + t.Run("update Dockerfile with new tag", func(t *testing.T) { + // Create a temporary old Dockerfile + oldDockerfilePath := "test_old_Dockerfile" + oldContent := "FROM quay.io/astronomer/astro-runtime:old-tag\n" + err := os.WriteFile(oldDockerfilePath, []byte(oldContent), 0o644) + assert.NoError(t, err) + defer os.Remove(oldDockerfilePath) + + // Define test data + newDockerfilePath := "test_new_Dockerfile" + newTag := "new-tag" + + // Call the function + err = upgradeDockerfile(oldDockerfilePath, newDockerfilePath, newTag, "") + defer os.Remove(newDockerfilePath) + + // Check for errors + assert.NoError(t, err) + + // Read the new Dockerfile and check its content + newContent, err := os.ReadFile(newDockerfilePath) + assert.NoError(t, err) + assert.Contains(t, string(newContent), "FROM quay.io/astronomer/astro-runtime:new-tag") + }) + + t.Run("update Dockerfile with new image", func(t *testing.T) { + // Create a temporary old Dockerfile + oldDockerfilePath := "test_old_Dockerfile" + oldContent := "FROM quay.io/astronomer/astro-runtime:old-tag\n" + err := os.WriteFile(oldDockerfilePath, []byte(oldContent), 0o644) + assert.NoError(t, err) + defer os.Remove(oldDockerfilePath) + + // Define test data + newDockerfilePath := "test_new_Dockerfile" + newImage := "new-image" + + // Call the function + err = upgradeDockerfile(oldDockerfilePath, newDockerfilePath, "", newImage) + defer os.Remove(newDockerfilePath) + + // Check for errors + assert.NoError(t, err) + + // Read the new Dockerfile and check its content + newContent, err := os.ReadFile(newDockerfilePath) + assert.NoError(t, err) + assert.Contains(t, string(newContent), "FROM new-image") + }) + + t.Run("update Dockerfile for ap-airflow with runtime version", func(t *testing.T) { + // Create a temporary old Dockerfile with a line matching the pattern + oldDockerfilePath := "test_old_Dockerfile" + oldContent := "FROM quay.io/astronomer/ap-airflow:old-tag" + err := os.WriteFile(oldDockerfilePath, []byte(oldContent), 0o644) + assert.NoError(t, err) + defer os.Remove(oldDockerfilePath) + + // Define test data + newDockerfilePath := "test_new_Dockerfile" + newTag := "5.0.0" + + // Call the function + err = upgradeDockerfile(oldDockerfilePath, newDockerfilePath, newTag, "") + defer os.Remove(newDockerfilePath) + + // Check for errors + assert.NoError(t, err) + + // Read the new Dockerfile and check its content + newContent, err := os.ReadFile(newDockerfilePath) + assert.NoError(t, err) + assert.Contains(t, string(newContent), "FROM quay.io/astronomer/astro-runtime:5.0.0\n") + }) + + t.Run("update Dockerfile for ap-airflow with non-runtime version", func(t *testing.T) { + // Create a temporary old Dockerfile with a line matching the pattern + oldDockerfilePath := "test_old_Dockerfile" + oldContent := "FROM quay.io/astronomer/ap-airflow:old-tag\n" + err := os.WriteFile(oldDockerfilePath, []byte(oldContent), 0o644) + assert.NoError(t, err) + defer os.Remove(oldDockerfilePath) + + // Define test data + newDockerfilePath := "test_new_Dockerfile" + newTag := "new-tag" + + // Call the function + err = upgradeDockerfile(oldDockerfilePath, newDockerfilePath, newTag, "") + defer os.Remove(newDockerfilePath) + + // Check for errors + assert.NoError(t, err) + + // Read the new Dockerfile and check its content + newContent, err := os.ReadFile(newDockerfilePath) + assert.NoError(t, err) + assert.Contains(t, string(newContent), "FROM quay.io/astronomer/ap-airflow:new-tag") + }) + + t.Run("error reading old Dockerfile", func(t *testing.T) { + // Define test data with an invalid path + oldDockerfilePath := "non_existent_Dockerfile" + newDockerfilePath := "test_new_Dockerfile" + newTag := "new-tag" + + // Call the function + err := upgradeDockerfile(oldDockerfilePath, newDockerfilePath, newTag, "") + defer os.Remove(newDockerfilePath) + + // Check for errors + assert.Error(t, err) + assert.Contains(t, err.Error(), "no such file or directory") + }) +} diff --git a/airflow/include/dagintegritytestdefault.py b/airflow/include/dagintegritytestdefault.py index 717d34433..89b0f4e52 100644 --- a/airflow/include/dagintegritytestdefault.py +++ b/airflow/include/dagintegritytestdefault.py @@ -99,7 +99,7 @@ def suppress_logging(namespace): def get_import_errors(): """ - Generate a tuple for import errors in the dag bag + Generate a tuple for import errors in the dag bag, and include DAGs without errors. """ with suppress_logging("airflow"): dag_bag = DagBag(include_examples=False) @@ -107,16 +107,30 @@ def get_import_errors(): def strip_path_prefix(path): return os.path.relpath(path, os.environ.get("AIRFLOW_HOME")) - # prepend "(None,None)" to ensure that a test object is always created even if it's a no op. - return [(None, None)] + [ - (strip_path_prefix(k), v.strip()) for k, v in dag_bag.import_errors.items() - ] + # Initialize an empty list to store the tuples + result = [] + + # Iterate over the items in import_errors + for k, v in dag_bag.import_errors.items(): + result.append((strip_path_prefix(k), v.strip())) + + # Check if there are DAGs without errors + for file_path in dag_bag.dags: + # Check if the file_path is not in import_errors, meaning no errors + if file_path not in dag_bag.import_errors: + result.append((strip_path_prefix(file_path), "No import errors")) + + return result @pytest.mark.parametrize( - "rel_path,rv", get_import_errors(), ids=[x[0] for x in get_import_errors()] + "rel_path, rv", get_import_errors(), ids=[x[0] for x in get_import_errors()] ) def test_file_imports(rel_path, rv): """Test for import errors on a file""" - if rel_path and rv: # Make sure our no op test doesn't raise an error + if rv != "No import errors": + # If rv is not "No import errors," consider it a failed test raise Exception(f"{rel_path} failed to import with message \n {rv}") + else: + # If rv is "No import errors," consider it a passed test + print(f"{rel_path} passed the import test") diff --git a/airflow/mocks/ContainerHandler.go b/airflow/mocks/ContainerHandler.go index 56cbe6c36..6182bd686 100644 --- a/airflow/mocks/ContainerHandler.go +++ b/airflow/mocks/ContainerHandler.go @@ -1,4 +1,4 @@ -// Code generated by mockery v2.20.2. DO NOT EDIT. +// Code generated by mockery v2.32.0. DO NOT EDIT. package mocks @@ -227,13 +227,12 @@ func (_m *ContainerHandler) UpgradeTest(runtimeVersion string, deploymentID stri return r0 } -type mockConstructorTestingTNewContainerHandler interface { +// NewContainerHandler creates a new instance of ContainerHandler. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations. +// The first argument is typically a *testing.T value. +func NewContainerHandler(t interface { mock.TestingT Cleanup(func()) -} - -// NewContainerHandler creates a new instance of ContainerHandler. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations. -func NewContainerHandler(t mockConstructorTestingTNewContainerHandler) *ContainerHandler { +}) *ContainerHandler { mock := &ContainerHandler{} mock.Mock.Test(t) diff --git a/airflow/mocks/DockerCLIClient.go b/airflow/mocks/DockerCLIClient.go index e40729cdd..575b37133 100644 --- a/airflow/mocks/DockerCLIClient.go +++ b/airflow/mocks/DockerCLIClient.go @@ -1,4 +1,4 @@ -// Code generated by mockery v2.20.2. DO NOT EDIT. +// Code generated by mockery v2.32.0. DO NOT EDIT. package mocks @@ -2642,13 +2642,12 @@ func (_m *DockerCLIClient) VolumesPrune(ctx context.Context, pruneFilter filters return r0, r1 } -type mockConstructorTestingTNewDockerCLIClient interface { +// NewDockerCLIClient creates a new instance of DockerCLIClient. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations. +// The first argument is typically a *testing.T value. +func NewDockerCLIClient(t interface { mock.TestingT Cleanup(func()) -} - -// NewDockerCLIClient creates a new instance of DockerCLIClient. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations. -func NewDockerCLIClient(t mockConstructorTestingTNewDockerCLIClient) *DockerCLIClient { +}) *DockerCLIClient { mock := &DockerCLIClient{} mock.Mock.Test(t) diff --git a/airflow/mocks/DockerComposeAPI.go b/airflow/mocks/DockerComposeAPI.go index 53c51c9df..71bf649bb 100644 --- a/airflow/mocks/DockerComposeAPI.go +++ b/airflow/mocks/DockerComposeAPI.go @@ -1,4 +1,4 @@ -// Code generated by mockery v2.20.2. DO NOT EDIT. +// Code generated by mockery v2.32.0. DO NOT EDIT. package mocks @@ -450,13 +450,12 @@ func (_m *DockerComposeAPI) Up(ctx context.Context, project *types.Project, opti return r0 } -type mockConstructorTestingTNewDockerComposeAPI interface { +// NewDockerComposeAPI creates a new instance of DockerComposeAPI. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations. +// The first argument is typically a *testing.T value. +func NewDockerComposeAPI(t interface { mock.TestingT Cleanup(func()) -} - -// NewDockerComposeAPI creates a new instance of DockerComposeAPI. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations. -func NewDockerComposeAPI(t mockConstructorTestingTNewDockerComposeAPI) *DockerComposeAPI { +}) *DockerComposeAPI { mock := &DockerComposeAPI{} mock.Mock.Test(t) diff --git a/airflow/mocks/DockerRegistryAPI.go b/airflow/mocks/DockerRegistryAPI.go index 1b4c4ba04..1981b4798 100644 --- a/airflow/mocks/DockerRegistryAPI.go +++ b/airflow/mocks/DockerRegistryAPI.go @@ -1,4 +1,4 @@ -// Code generated by mockery v2.20.2. DO NOT EDIT. +// Code generated by mockery v2.32.0. DO NOT EDIT. package mocks @@ -2588,13 +2588,12 @@ func (_m *DockerRegistryAPI) VolumesPrune(ctx context.Context, pruneFilter filte return r0, r1 } -type mockConstructorTestingTNewDockerRegistryAPI interface { +// NewDockerRegistryAPI creates a new instance of DockerRegistryAPI. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations. +// The first argument is typically a *testing.T value. +func NewDockerRegistryAPI(t interface { mock.TestingT Cleanup(func()) -} - -// NewDockerRegistryAPI creates a new instance of DockerRegistryAPI. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations. -func NewDockerRegistryAPI(t mockConstructorTestingTNewDockerRegistryAPI) *DockerRegistryAPI { +}) *DockerRegistryAPI { mock := &DockerRegistryAPI{} mock.Mock.Test(t) diff --git a/airflow/mocks/ImageHandler.go b/airflow/mocks/ImageHandler.go index a70b41876..da33fafd2 100644 --- a/airflow/mocks/ImageHandler.go +++ b/airflow/mocks/ImageHandler.go @@ -1,4 +1,4 @@ -// Code generated by mockery v2.20.2. DO NOT EDIT. +// Code generated by mockery v2.32.0. DO NOT EDIT. package mocks @@ -64,6 +64,20 @@ func (_m *ImageHandler) CreatePipFreeze(altImageName string, pipFreezeFile strin return r0 } +// DoesImageExist provides a mock function with given fields: image +func (_m *ImageHandler) DoesImageExist(image string) error { + ret := _m.Called(image) + + var r0 error + if rf, ok := ret.Get(0).(func(string) error); ok { + r0 = rf(image) + } else { + r0 = ret.Error(0) + } + + return r0 +} + // GetLabel provides a mock function with given fields: altImageName, labelName func (_m *ImageHandler) GetLabel(altImageName string, labelName string) (string, error) { ret := _m.Called(altImageName, labelName) @@ -194,13 +208,12 @@ func (_m *ImageHandler) TagLocalImage(localImage string) error { return r0 } -type mockConstructorTestingTNewImageHandler interface { +// NewImageHandler creates a new instance of ImageHandler. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations. +// The first argument is typically a *testing.T value. +func NewImageHandler(t interface { mock.TestingT Cleanup(func()) -} - -// NewImageHandler creates a new instance of ImageHandler. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations. -func NewImageHandler(t mockConstructorTestingTNewImageHandler) *ImageHandler { +}) *ImageHandler { mock := &ImageHandler{} mock.Mock.Test(t) diff --git a/airflow/mocks/RegistryHandler.go b/airflow/mocks/RegistryHandler.go index 15caacd7f..49eb516a9 100644 --- a/airflow/mocks/RegistryHandler.go +++ b/airflow/mocks/RegistryHandler.go @@ -1,4 +1,4 @@ -// Code generated by mockery v2.20.2. DO NOT EDIT. +// Code generated by mockery v2.32.0. DO NOT EDIT. package mocks @@ -23,13 +23,12 @@ func (_m *RegistryHandler) Login(username string, token string) error { return r0 } -type mockConstructorTestingTNewRegistryHandler interface { +// NewRegistryHandler creates a new instance of RegistryHandler. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations. +// The first argument is typically a *testing.T value. +func NewRegistryHandler(t interface { mock.TestingT Cleanup(func()) -} - -// NewRegistryHandler creates a new instance of RegistryHandler. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations. -func NewRegistryHandler(t mockConstructorTestingTNewRegistryHandler) *RegistryHandler { +}) *RegistryHandler { mock := &RegistryHandler{} mock.Mock.Test(t) diff --git a/astro-client-core/mocks/client.go b/astro-client-core/mocks/client.go index 231f8e898..4c2743611 100644 --- a/astro-client-core/mocks/client.go +++ b/astro-client-core/mocks/client.go @@ -1,4 +1,4 @@ -// Code generated by mockery v2.20.2. DO NOT EDIT. +// Code generated by mockery v2.32.0. DO NOT EDIT. package astrocore_mocks @@ -3845,13 +3845,12 @@ func (_m *ClientWithResponsesInterface) VerifyManagedDomainWithResponse(ctx cont return r0, r1 } -type mockConstructorTestingTNewClientWithResponsesInterface interface { +// NewClientWithResponsesInterface creates a new instance of ClientWithResponsesInterface. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations. +// The first argument is typically a *testing.T value. +func NewClientWithResponsesInterface(t interface { mock.TestingT Cleanup(func()) -} - -// NewClientWithResponsesInterface creates a new instance of ClientWithResponsesInterface. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations. -func NewClientWithResponsesInterface(t mockConstructorTestingTNewClientWithResponsesInterface) *ClientWithResponsesInterface { +}) *ClientWithResponsesInterface { mock := &ClientWithResponsesInterface{} mock.Mock.Test(t) diff --git a/astro-client/mocks/Client.go b/astro-client/mocks/Client.go index c93df6fd1..90680be96 100644 --- a/astro-client/mocks/Client.go +++ b/astro-client/mocks/Client.go @@ -1,4 +1,4 @@ -// Code generated by mockery v2.20.2. DO NOT EDIT. +// Code generated by mockery v2.32.0. DO NOT EDIT. package astro_mocks @@ -385,13 +385,12 @@ func (_m *Client) UpdateDeployment(input *astro.UpdateDeploymentInput) (astro.De return r0, r1 } -type mockConstructorTestingTNewClient interface { +// NewClient creates a new instance of Client. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations. +// The first argument is typically a *testing.T value. +func NewClient(t interface { mock.TestingT Cleanup(func()) -} - -// NewClient creates a new instance of Client. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations. -func NewClient(t mockConstructorTestingTNewClient) *Client { +}) *Client { mock := &Client{} mock.Mock.Test(t) diff --git a/cmd/airflow.go b/cmd/airflow.go index 2642e869b..2067ebd82 100644 --- a/cmd/airflow.go +++ b/cmd/airflow.go @@ -549,25 +549,44 @@ func airflowInit(cmd *cobra.Command, args []string) error { return nil } -func airflowUpgradeTest(cmd *cobra.Command, astroClient astro.Client) error { +func airflowUpgradeTest(cmd *cobra.Command, astroClient astro.Client) error { //nolint:gocognit // Validate runtimeVersion and airflowVersion if airflowVersion != "" && runtimeVersion != "" { - return errInvalidBothAirflowAndRuntimeVersions + return errInvalidBothAirflowAndRuntimeVersionsUpgrade } - // If user provides a runtime version, use it, otherwise retrieve the latest one (matching Airflow Version if provided) - var err error - defaultImageTag := runtimeVersion - if defaultImageTag == "" { - httpClient := airflowversions.NewClient(httputil.NewHTTPClient(), useAstronomerCertified) - defaultImageTag = prepareDefaultAirflowImageTag(airflowVersion, httpClient) + // error if both custom image and deployment id is used + if deploymentID != "" && customImageName != "" { + return errInvalidBothDeploymentIDandCustomImage + } + if airflowVersion != "" && deploymentID != "" { + return errInvalidBothDeploymentIDandVersion + } + if runtimeVersion != "" && deploymentID != "" { + return errInvalidBothDeploymentIDandVersion + } + if runtimeVersion != "" && customImageName != "" { + return errInvalidBothCustomImageandVersion + } + if airflowVersion != "" && customImageName != "" { + return errInvalidBothCustomImageandVersion } defaultImageName := airflow.AstroRuntimeImageName - if useAstronomerCertified { - defaultImageName = airflow.AstronomerCertifiedImageName - fmt.Printf("Testing an upgrade to Astronomer Certified Airflow version %s\n\n", defaultImageTag) + defaultImageTag := runtimeVersion + if customImageName != "" { + fmt.Printf("Testing an upgrade to custom Airflow image: %s\n", customImageName) } else { - fmt.Printf("Testing an upgrade to Astro Runtime %s\n", defaultImageTag) + // If user provides a runtime version, use it, otherwise retrieve the latest one (matching Airflow Version if provided + if defaultImageTag == "" { + httpClient := airflowversions.NewClient(httputil.NewHTTPClient(), useAstronomerCertified) + defaultImageTag = prepareDefaultAirflowImageTag(airflowVersion, httpClient) + } + if useAstronomerCertified { + defaultImageName = airflow.AstronomerCertifiedImageName + fmt.Printf("Testing an upgrade to Astronomer Certified Airflow version %s\n\n", defaultImageTag) + } else { + fmt.Printf("Testing an upgrade to Astro Runtime %s\n", defaultImageTag) + } } // Silence Usage as we have now validated command input diff --git a/cmd/airflow_test.go b/cmd/airflow_test.go index 91d20ff65..b1d01e9de 100644 --- a/cmd/airflow_test.go +++ b/cmd/airflow_test.go @@ -483,6 +483,66 @@ func TestAirflowUpgradeTest(t *testing.T) { err := airflowUpgradeTest(cmd, nil) assert.ErrorIs(t, err, errMock) }) + + t.Run("Both airflow and runtime version used", func(t *testing.T) { + cmd := newAirflowUpgradeTestCmd(nil) + + airflowVersion = "something" + runtimeVersion = "something" + + err := airflowUpgradeTest(cmd, nil) + assert.ErrorIs(t, err, errInvalidBothAirflowAndRuntimeVersionsUpgrade) + }) + + t.Run("Both custom image and deployment id used", func(t *testing.T) { + cmd := newAirflowUpgradeTestCmd(nil) + + deploymentID = "something" + customImageName = "something" + + err := airflowUpgradeTest(cmd, nil) + assert.ErrorIs(t, err, errInvalidBothDeploymentIDandCustomImage) + }) + + t.Run("Both airflow version and deployment id used", func(t *testing.T) { + cmd := newAirflowUpgradeTestCmd(nil) + + deploymentID = "something" + airflowVersion = "something" + + err := airflowUpgradeTest(cmd, nil) + assert.ErrorIs(t, err, errInvalidBothDeploymentIDandVersion) + }) + + t.Run("Both runtime version and deployment id used", func(t *testing.T) { + cmd := newAirflowUpgradeTestCmd(nil) + + deploymentID = "something" + runtimeVersion = "something" + + err := airflowUpgradeTest(cmd, nil) + assert.ErrorIs(t, err, errInvalidBothDeploymentIDandVersion) + }) + + t.Run("Both runtime version and custom image used", func(t *testing.T) { + cmd := newAirflowUpgradeTestCmd(nil) + + customImageName = "something" + runtimeVersion = "something" + + err := airflowUpgradeTest(cmd, nil) + assert.ErrorIs(t, err, errInvalidBothCustomImageandVersion) + }) + + t.Run("Both airflow version and custom image used", func(t *testing.T) { + cmd := newAirflowUpgradeTestCmd(nil) + + customImageName = "something" + airflowVersion = "something" + + err := airflowUpgradeTest(cmd, nil) + assert.ErrorIs(t, err, errInvalidBothCustomImageandVersion) + }) } func TestAirflowRun(t *testing.T) { diff --git a/cmd/errors.go b/cmd/errors.go index 5cb007023..2d8d83800 100644 --- a/cmd/errors.go +++ b/cmd/errors.go @@ -5,7 +5,11 @@ import ( ) var ( - errInvalidBothAirflowAndRuntimeVersions = errors.New("You provided both a runtime version and an Airflow version. You have to provide only one of these to initialize your project.") //nolint + errInvalidBothAirflowAndRuntimeVersions = errors.New("You provided both a runtime version and an Airflow version. You have to provide only one of these to initialize your project.") //nolint + errInvalidBothAirflowAndRuntimeVersionsUpgrade = errors.New("You provided both a runtime version and an Airflow version. You have to provide only one of these to upgrade.") //nolint + errInvalidBothDeploymentIDandCustomImage = errors.New("You provided both a Deployment ID and a Custom image. You have to provide only one of these to upgrade.") //nolint + errInvalidBothDeploymentIDandVersion = errors.New("You provided both a Deployment ID and a version. You have to provide only one of these to upgrade.") //nolint + errInvalidBothCustomImageandVersion = errors.New("You provided both a Custom image and a version. You have to provide only one of these to upgrade.") //nolint errConfigProjectName = errors.New("project name is invalid") errProjectNameSpaces = errors.New("this project name is invalid, a project name cannot contain spaces. Try using '-' instead") diff --git a/houston/mocks/ClientInterface.go b/houston/mocks/ClientInterface.go index 03acfe3c5..2e39a38b0 100644 --- a/houston/mocks/ClientInterface.go +++ b/houston/mocks/ClientInterface.go @@ -1,4 +1,4 @@ -// Code generated by mockery v2.20.2. DO NOT EDIT. +// Code generated by mockery v2.32.0. DO NOT EDIT. package houston_mocks @@ -1428,13 +1428,12 @@ func (_m *ClientInterface) ValidateWorkspaceID(workspaceID string) (*houston.Wor return r0, r1 } -type mockConstructorTestingTNewClientInterface interface { +// NewClientInterface creates a new instance of ClientInterface. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations. +// The first argument is typically a *testing.T value. +func NewClientInterface(t interface { mock.TestingT Cleanup(func()) -} - -// NewClientInterface creates a new instance of ClientInterface. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations. -func NewClientInterface(t mockConstructorTestingTNewClientInterface) *ClientInterface { +}) *ClientInterface { mock := &ClientInterface{} mock.Mock.Test(t) diff --git a/pkg/azure/mocks/Azure.go b/pkg/azure/mocks/Azure.go index d6bb35a5f..66853a879 100644 --- a/pkg/azure/mocks/Azure.go +++ b/pkg/azure/mocks/Azure.go @@ -1,4 +1,4 @@ -// Code generated by mockery v2.20.2. DO NOT EDIT. +// Code generated by mockery v2.32.0. DO NOT EDIT. package azure_mocks @@ -37,13 +37,12 @@ func (_m *Azure) Upload(sasLink string, dagFileReader io.Reader) (string, error) return r0, r1 } -type mockConstructorTestingTNewAzure interface { +// NewAzure creates a new instance of Azure. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations. +// The first argument is typically a *testing.T value. +func NewAzure(t interface { mock.TestingT Cleanup(func()) -} - -// NewAzure creates a new instance of Azure. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations. -func NewAzure(t mockConstructorTestingTNewAzure) *Azure { +}) *Azure { mock := &Azure{} mock.Mock.Test(t)