Upgrade command (#1314)

* dependency conflict check * add dependency compare * add all tests * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * update tests * Fixing tests * fix test * fix lint * fix lint * fix lint * fix lint * fix lint * fix lint * empty commit * fix lint * fix lint * fix lint * fix lint * fix lint * fix lint * fix lint * empty commit * fix lint * add tests * move function * add certified flag * add test * add test * add test * fix test * fix test * fix test * fix test * fix test * fix test * add test * add test * add test * fix test * merge with main * merge with main * Apply suggestions from code review Co-authored-by: Neel Dalsania <[email protected]> Co-authored-by: kushalmalani <[email protected]> * fix test * fix test * Apply suggestions from code review Co-authored-by: Jake Witz <[email protected]> * update * update copy * update copy * move to test files * Apply suggestions from code review Co-authored-by: Jake Witz <[email protected]> * Update cmd/airflow.go Co-authored-by: Jake Witz <[email protected]> * Update cmd/airflow.go Co-authored-by: Jake Witz <[email protected]> * fix test * add custom image * fix command * fix test * fix dockerfile upgrade * remove init examples * fix command * Update airflow/docker_image.go Co-authored-by: Neel Dalsania <[email protected]> * fix lint * fix lint * fix lint * fix lint * fix lint --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Kushal Malani <[email protected]> Co-authored-by: Neel Dalsania <[email protected]> Co-authored-by: Jake Witz <[email protected]>
astronomer · Aug 3, 2023 · 595721e · 595721e
1 parent beca3f0
commit 595721e
Show file tree

Hide file tree

Showing 28 changed files with 2,146 additions and 187 deletions.
diff --git a/airflow-client/mocks/Client.go b/airflow-client/mocks/Client.go
diff --git a/airflow/airflow.go b/airflow/airflow.go
@@ -34,6 +34,9 @@ var (
 	//go:embed include/dockerfile
 	Dockerfile string
 
+	//go:embed include/test-conflicts.dockerfile
+	testConflictsDockerfile string
+
 	//go:embed include/dockerignore
 	Dockerignore string
 
@@ -131,6 +134,18 @@ func Init(path, airflowImageName, airflowImageTag string) error {
 	return nil
 }
 
+func initConflictTest(path, airflowImageName, airflowImageTag string) error {
+	// Map of files to create
+	files := map[string]string{
+		"conflict-check.Dockerfile": fmt.Sprintf(testConflictsDockerfile, airflowImageName, airflowImageTag),
+	}
+	// Initialize files
+	if err := initFiles(path, files); err != nil {
+		return errors.Wrap(err, "failed to create upgrade check files")
+	}
+	return nil
+}
+
 // repositoryName creates an airflow repository name
 func repositoryName(name string) string {
 	return fmt.Sprintf("%s/%s", name, componentName)

diff --git a/airflow/airflow_test.go b/airflow/airflow_test.go
@@ -87,3 +87,23 @@ func TestInit(t *testing.T) {
 		assert.True(t, exist)
 	}
 }
+
+func TestInitConflictTest(t *testing.T) {
+	tmpDir, err := os.MkdirTemp("", "temp")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer os.RemoveAll(tmpDir)
+
+	err = initConflictTest(tmpDir, "astro-runtime", "test")
+	assert.NoError(t, err)
+
+	expectedFiles := []string{
+		"conflict-check.Dockerfile",
+	}
+	for _, file := range expectedFiles {
+		exist, err := fileutil.Exists(filepath.Join(tmpDir, file), nil)
+		assert.NoError(t, err)
+		assert.True(t, exist)
+	}
+}
diff --git a/airflow/container.go b/airflow/container.go
@@ -10,6 +10,7 @@ import (
 	"time"
 
 	"github.com/astronomer/astro-cli/airflow/types"
+	"github.com/astronomer/astro-cli/astro-client"
 	"github.com/astronomer/astro-cli/config"
 	"github.com/astronomer/astro-cli/pkg/fileutil"
 	"github.com/astronomer/astro-cli/pkg/util"
@@ -33,6 +34,7 @@ type ContainerHandler interface {
 	ComposeExport(settingsFile, composeFile string) error
 	Pytest(pytestFile, customImageName, deployImageName, pytestArgsString string) (string, error)
 	Parse(customImageName, deployImageName string) error
+	UpgradeTest(runtimeVersion, deploymentID, newImageName, customImageName string, dependencyTest, versionTest, dagTest bool, client astro.Client) error
 }
 
 // RegistryHandler defines methods require to handle all operations with registry
@@ -42,13 +44,16 @@ type RegistryHandler interface {
 
 // ImageHandler defines methods require to handle all operations on/for container images
 type ImageHandler interface {
-	Build(config types.ImageBuildConfig) error
+	Build(dockerfile string, config types.ImageBuildConfig) error
 	Push(registry, username, token, remoteImage string) error
-	GetLabel(labelName string) (string, error)
+	Pull(registry, username, token, remoteImage string) error
+	GetLabel(altImageName, labelName string) (string, error)
 	ListLabels() (map[string]string, error)
 	TagLocalImage(localImage string) error
 	Run(dagID, envFile, settingsFile, containerName, dagFile, executionDate string, taskLogs bool) error
-	Pytest(pytestFile, airflowHome, envFile string, pytestArgs []string, config types.ImageBuildConfig) (string, error)
+	Pytest(pytestFile, airflowHome, envFile, testHomeDirectory string, pytestArgs []string, htmlReport bool, config types.ImageBuildConfig) (string, error)
+	ConflictTest(workingDirectory, testHomeDirectory string, buildConfig types.ImageBuildConfig) (string, error)
+	CreatePipFreeze(altImageName, pipFreezeFile string) error
 }
 
 type DockerComposeAPI interface {