Merge pull request #413 from astronomer/hotfix/astro-dev-init-set-ari…

…flow-version

fix astro dev init -v

cmd/airflow.go

@@ -276,17 +276,9 @@ func airflowInit(cmd *cobra.Command, args []string, client *houston.Client, out
 		projectName = strings.Replace(strcase.ToSnake(projectDirectory), "_", "-", -1)
 	}
 	httpClient := airflowversions.NewClient(httputil.NewHTTPClient())
-	defaultImageTag, _ := airflowversions.GetDefaultImageTag(httpClient, "")
-
-	// TODO: @andriisoldatenko rethink or remove this logic
-	// acceptableAirflowVersions := wsResp.Data.DeploymentConfig.AirflowVersions
-	// if airflowVersion != "" && !acceptableVersion(airflowVersion, acceptableAirflowVersions) {
-	//  	return errors.Errorf(messages.ERROR_INVALID_AIRFLOW_VERSION, strings.Join(acceptableAirflowVersions, ", "))
-	// }
-
-	if len(defaultImageTag) == 0 {
-		defaultImageTag = "2.0.0-buster-onbuild"
-		fmt.Printf("Initializing Airflow project\nNot connected to Astronomer, pulling Airflow development files from %s\n", defaultImageTag)
+	defaultImageTag, err := prepareDefaultAirflowImageTag(airflowVersion, httpClient, client, out)
+	if err != nil {
+		return err
 	}
 
 	emtpyDir := fileutil.IsEmptyDir(config.WorkingPath)
@@ -309,7 +301,7 @@ func airflowInit(cmd *cobra.Command, args []string, client *houston.Client, out
 	cmd.SilenceUsage = true
 
 	// Execute method
-	err := airflow.Init(config.WorkingPath, defaultImageTag)
+	err = airflow.Init(config.WorkingPath, defaultImageTag)
 	if err != nil {
 		return err
 	}

cmd/utils.go

@@ -0,0 +1,42 @@
+package cmd
+
+import (
+	"fmt"
+	"io"
+	"strings"
+
+	airflowversions "github.com/astronomer/astro-cli/airflow_versions"
+	"github.com/astronomer/astro-cli/houston"
+	"github.com/astronomer/astro-cli/messages"
+	"github.com/pkg/errors"
+)
+
+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,
+	}
+
+	wsResp, err := r.DoWithClient(houstonClient)
+
+	if err == nil {
+		acceptableAirflowVersions := wsResp.Data.DeploymentConfig.AirflowVersions
+		if airflowVersion != "" && !acceptableVersion(airflowVersion, acceptableAirflowVersions) {
+			return "", errors.Errorf(messages.ERROR_INVALID_AIRFLOW_VERSION, strings.Join(acceptableAirflowVersions, ", "))
+		}
+		if airflowVersion == "" {
+			defaultImageTag = ""
+		} else {
+			defaultImageTag = fmt.Sprintf("%s-buster-onbuild", airflowVersion)
+		}
+	} else if airflowVersion != "" {
+		return "", errors.New("An Error occurred: Are you authenticated? If not - you can't use the --airflow-version option")
+	}
+
+	if len(defaultImageTag) == 0 {
+		defaultImageTag = "2.0.0-buster-onbuild"
+		fmt.Fprintf(out, "Initializing Airflow project\nNot connected to Astronomer, pulling Airflow development files from %s\n", defaultImageTag)
+	}
+	return defaultImageTag, nil
+}

cmd/utils_test.go

@@ -0,0 +1,144 @@
+package cmd
+
+import (
+	"bytes"
+	"io/ioutil"
+	"net/http"
+	"testing"
+
+	airflowversions "github.com/astronomer/astro-cli/airflow_versions"
+	"github.com/astronomer/astro-cli/houston"
+	testUtil "github.com/astronomer/astro-cli/pkg/testing"
+	"github.com/stretchr/testify/assert"
+)
+
+func Test_prepareDefaultAirflowImageTag(t *testing.T) {
+	testUtil.InitTestConfig()
+
+	// prepare fake response from updates.astronomer.io
+	okResponse := `{
+  "version": "1.0",
+  "available_releases": [
+    {
+      "version": "1.10.5",
+      "level": "new_feature",
+      "url": "https://github.com/astronomer/airflow/releases/tag/1.10.5-11",
+      "release_date": "2020-10-05T20:03:00+00:00",
+      "tags": [
+        "1.10.5-alpine3.10-onbuild",
+        "1.10.5-buster-onbuild",
+        "1.10.5-alpine3.10",
+        "1.10.5-buster"
+      ],
+      "channel": "stable"
+    }
+  ]
+}`
+	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: "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: "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 {
+		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()
+
+	// prepare fake response from updates.astronomer.io
+	okResponse := `{
+  "version": "1.0",
+  "available_releases": [
+    {
+      "version": "1.10.5",
+      "level": "new_feature",
+      "url": "https://github.com/astronomer/airflow/releases/tag/1.10.5-11",
+      "release_date": "2020-10-05T20:03:00+00:00",
+      "tags": [
+        "1.10.5-alpine3.10-onbuild",
+        "1.10.5-buster-onbuild",
+        "1.10.5-alpine3.10",
+        "1.10.5-buster"
+      ],
+      "channel": "stable"
+    }
+  ]
+}`
+	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
+	emptyResp := ``
+	houstonClient := testUtil.NewTestClient(func(req *http.Request) *http.Response {
+		return &http.Response{
+			StatusCode: 400,
+			Body:       ioutil.NopCloser(bytes.NewBufferString(emptyResp)),
+			Header:     make(http.Header),
+		}
+	})
+	api := houston.NewHoustonClient(houstonClient)
+
+	output := new(bytes.Buffer)
+
+	defaultTag, err := prepareDefaultAirflowImageTag("2.0.2", httpClient, api, output)
+	assert.Error(t, err, "An Error occurred: Are you authenticated? If not - you can't use the --airflow-version option")
+	assert.Equal(t, "", defaultTag)
+}