fixed default airflow image tag behaviour (#450)

* fixed default airflow image tag behaviour

* added test to cover new default airflow image tag logic

* fixed the output string for selecting default airflow image

* added code review changes

airflow_versions/airflow_versions.go

@@ -3,8 +3,11 @@ package airflowversions
 import (
 	"fmt"
 	"sort"
+	"strings"
 )
 
+var tagPrefixOrder = []string{"buster-onbuild", "onbuild", "buster"}
+
 // GetDefaultImageTag returns default airflow image tag
 func GetDefaultImageTag(httpClient *Client, airflowVersion string) (string, error) {
 	r := Request{}
@@ -14,15 +17,39 @@ func GetDefaultImageTag(httpClient *Client, airflowVersion string) (string, erro
 		return "", err
 	}
 
+	availableTags := []string{}
 	vs := make(AirflowVersions, len(resp.AvailableReleases))
 	for i, r := range resp.AvailableReleases {
+		if r.Version == airflowVersion {
+			availableTags = r.Tags
+			break
+		}
 		v, err := NewAirflowVersion(r.Version, r.Tags)
 		if err == nil {
 			vs[i] = v
 		}
 	}
 
-	sort.Sort(vs)
-	maxAvailableVersion := vs[len(vs)-1]
-	return fmt.Sprintf("%s-buster-onbuild", maxAvailableVersion.Coerce()), nil
+	var selectedVersion *AirflowVersion
+	if airflowVersion == "" && len(vs) != 0 {
+		sort.Sort(vs)
+		selectedVersion = vs[len(vs)-1]
+		availableTags = selectedVersion.tags
+	} else {
+		selectedVersion, err = NewAirflowVersion(airflowVersion, availableTags)
+		if err != nil {
+			return "", err
+		}
+	}
+
+	for tagIndex := range tagPrefixOrder {
+		for idx := range availableTags {
+			if strings.HasPrefix(availableTags[idx], selectedVersion.Coerce()) && strings.HasSuffix(availableTags[idx], tagPrefixOrder[tagIndex]) {
+				return fmt.Sprintf("%s-%s", selectedVersion.Coerce(), tagPrefixOrder[tagIndex]), nil
+			}
+		}
+	}
+
+	// case when airflowVersion requested is not present in certified astronomer endpoint, but is valid version as per Houston configuration
+	return airflowVersion, nil
 }

airflow_versions/airflow_versions_test.go

@@ -8,6 +8,8 @@ import (
 
 	testUtil "github.com/astronomer/astro-cli/pkg/testing"
 	"github.com/stretchr/testify/assert"
+
+	"github.com/Masterminds/semver"
 )
 
 func TestGetDefaultImageTag(t *testing.T) {
@@ -53,6 +55,14 @@ func TestGetDefaultImageTag(t *testing.T) {
         "1.10.4-11-buster"
       ],
       "channel": "stable"
+    },
+    {
+      "version": "2.2.0",
+      "level": "new_feature",
+      "url": "https://github.com/astronomer/airflow/releases/tag/v2.2.0%2Bastro.2",
+      "release_date": "2021-10-14T12:46:00+00:00",
+      "tags": ["2.2.0", "2.2.0-onbuild"],
+      "channel": "stable"
     }
   ]
 }`
@@ -65,9 +75,28 @@ func TestGetDefaultImageTag(t *testing.T) {
 	})
 	httpClient := NewClient(client)
 
-	defaultImageTag, err := GetDefaultImageTag(httpClient, "")
-	assert.NoError(t, err)
-	assert.Equal(t, "1.10.5-buster-onbuild", defaultImageTag)
+	tests := []struct {
+		airflowVersion string
+		output         string
+		err            error
+	}{
+		{airflowVersion: "", output: "2.2.0-onbuild", err: nil},
+		{airflowVersion: "1.10.5", output: "1.10.5-buster-onbuild", err: nil},
+		{airflowVersion: "1.10.4-rc.1", output: "1.10.4-rc.1", err: nil},
+		{airflowVersion: "2.2.1", output: "2.2.1", err: nil},
+		{airflowVersion: "2.2.0", output: "2.2.0-onbuild", err: nil},
+		{airflowVersion: "2.2.x", output: "", err: semver.ErrInvalidSemVer},
+	}
+
+	for _, tt := range tests {
+		defaultImageTag, err := GetDefaultImageTag(httpClient, tt.airflowVersion)
+		if tt.err == nil {
+			assert.NoError(t, err)
+		} else {
+			assert.EqualError(t, err, tt.err.Error())
+		}
+		assert.Equal(t, tt.output, defaultImageTag)
+	}
 }
 
 func TestGetDefaultImageTagError(t *testing.T) {

cmd/utils.go

@@ -12,8 +12,6 @@ import (
 )
 
 func prepareDefaultAirflowImageTag(airflowVersion string, httpClient *airflowversions.Client, houstonClient *houston.Client, out io.Writer) (string, error) {
-	defaultImageTag, _ := airflowversions.GetDefaultImageTag(httpClient, "")
-
 	r := houston.Request{
 		Query: houston.DeploymentInfoRequest,
 	}
@@ -25,23 +23,20 @@ func prepareDefaultAirflowImageTag(airflowVersion string, httpClient *airflowver
 		if airflowVersion != "" && !acceptableVersion(airflowVersion, acceptableAirflowVersions) {
 			return "", errors.Errorf(messages.ErrInvalidAirflowVersion, strings.Join(acceptableAirflowVersions, ", "))
 		}
-		if airflowVersion == "" {
-			defaultImageTag = ""
-		} else {
-			defaultImageTag = fmt.Sprintf("%s-buster-onbuild", airflowVersion)
-		}
 	} else if airflowVersion != "" {
 		switch t := err; t {
-		default:
-			return "", err
 		case houston.ErrVerboseInaptPermissions:
 			return "", errors.New("the --airflow-version flag is not supported if you're not authenticated to Astronomer. Please authenticate and try again")
+		default:
+			return "", err
 		}
 	}
 
+	defaultImageTag, _ := airflowversions.GetDefaultImageTag(httpClient, airflowVersion)
+
 	if defaultImageTag == "" {
 		defaultImageTag = "2.0.0-buster-onbuild"
-		fmt.Fprintf(out, "Initializing Airflow project\nNot connected to Astronomer, pulling Airflow development files from %s\n", defaultImageTag)
+		fmt.Fprintf(out, "Initializing Airflow project, pulling Airflow development files from %s\n", defaultImageTag)
 	}
 	return defaultImageTag, nil
 }

cmd/utils_test.go

@@ -78,10 +78,10 @@ func Test_prepareDefaultAirflowImageTag(t *testing.T) {
 		expectedImageTag string
 		expectedError    string
 	}{
-		{airflowVersion: "1.10.14", expectedImageTag: "1.10.14-buster-onbuild", expectedError: ""},
-		{airflowVersion: "1.10.15", expectedImageTag: "1.10.15-buster-onbuild", expectedError: ""},
-		{airflowVersion: "", expectedImageTag: "2.0.0-buster-onbuild", expectedError: ""},
-		{airflowVersion: "2.0.2", expectedImageTag: "2.0.2-buster-onbuild", expectedError: ""},
+		{airflowVersion: "1.10.14", expectedImageTag: "1.10.14", expectedError: ""},
+		{airflowVersion: "1.10.15", expectedImageTag: "1.10.15", expectedError: ""},
+		{airflowVersion: "", expectedImageTag: "1.10.5-buster-onbuild", expectedError: ""},
+		{airflowVersion: "2.0.2", expectedImageTag: "2.0.2", expectedError: ""},
 		{airflowVersion: "9.9.9", expectedImageTag: "", expectedError: "Unsupported Airflow Version specified. Please choose from: 2.1.0, 2.0.2, 2.0.0, 1.10.15, 1.10.14, 1.10.12, 1.10.10, 1.10.7, 1.10.5 \n"},
 	}
 	for _, tt := range myTests {
@@ -95,6 +95,70 @@ func Test_prepareDefaultAirflowImageTag(t *testing.T) {
 	}
 }
 
+func Test_fallbackDefaultAirflowImageTag(t *testing.T) {
+	testUtil.InitTestConfig()
+
+	// prepare fake response from updates.astronomer.io
+	okResponse := `{
+  "version": "1.0",
+  "available_releases": []
+}`
+	client := testUtil.NewTestClient(func(req *http.Request) *http.Response {
+		return &http.Response{
+			StatusCode: 200,
+			Body:       ioutil.NopCloser(bytes.NewBufferString(okResponse)),
+			Header:     make(http.Header),
+		}
+	})
+	httpClient := airflowversions.NewClient(client)
+
+	// prepare fake response from houston
+	ok := `{
+	"data": {
+		"deploymentConfig": {
+				"airflowVersions": [
+					"2.1.0",
+					"2.0.2",
+					"2.0.0",
+					"1.10.15",
+					"1.10.14",
+					"1.10.12",
+					"1.10.10",
+					"1.10.7",
+					"1.10.5"
+				]
+			}
+		}
+	}`
+	houstonClient := testUtil.NewTestClient(func(req *http.Request) *http.Response {
+		return &http.Response{
+			StatusCode: 200,
+			Body:       ioutil.NopCloser(bytes.NewBufferString(ok)),
+			Header:     make(http.Header),
+		}
+	})
+	api := houston.NewHoustonClient(houstonClient)
+
+	output := new(bytes.Buffer)
+
+	myTests := []struct {
+		airflowVersion   string
+		expectedImageTag string
+		expectedError    string
+	}{
+		{airflowVersion: "", expectedImageTag: "2.0.0-buster-onbuild", expectedError: ""},
+	}
+	for _, tt := range myTests {
+		defaultTag, err := prepareDefaultAirflowImageTag(tt.airflowVersion, httpClient, api, output)
+		if tt.expectedError != "" {
+			assert.EqualError(t, err, tt.expectedError)
+		} else {
+			assert.NoError(t, err)
+		}
+		assert.Equal(t, tt.expectedImageTag, defaultTag)
+	}
+}
+
 func Test_prepareDefaultAirflowImageTagHoustonBadRequest(t *testing.T) {
 	testUtil.InitTestConfig()
 	mockErrorResponse := `An error occurred`