changes for python file entry point detection

cmd/create_test.go

@@ -15,6 +15,7 @@ import (
 	"github.com/Azure/draft/pkg/config"
 	"github.com/Azure/draft/pkg/languages"
 	"github.com/Azure/draft/pkg/linguist"
+	"github.com/Azure/draft/pkg/reporeader"
 	"github.com/Azure/draft/pkg/templatewriter/writers"
 	"github.com/Azure/draft/template"
 )
@@ -53,7 +54,6 @@ func TestRun(t *testing.T) {
 	assert.False(t, lowerLang == "")
 	assert.True(t, err == nil)
 	err = mockCC.generateDockerfile(detectedLang, lowerLang)
-	println(err)
 	assert.True(t, err == nil)
 
 	//Write back old Dockerfile
@@ -105,6 +105,43 @@ func TestRun(t *testing.T) {
 	}
 }
 
+func TestRunCreateDockerfileWithRepoReader(t *testing.T) {
+	mockCC := &createCmd{}
+	mockCC.repoReader = &reporeader.TestRepoReader{Files: map[string][]byte{
+		"foo.py":  []byte("print('Hello World')"),
+		"main.py": []byte("print('Hello World')"),
+	}}
+	mockCC.createConfig = &CreateConfig{}
+	mockCC.createConfig.LanguageType = "python"
+	mockCC.createConfig.LanguageVariables = []UserInputs{}
+	mockPortInput := UserInputs{Name: "PORT", Value: "8080"}
+	mockCC.createConfig.LanguageVariables = append(mockCC.createConfig.LanguageVariables, mockPortInput)
+	mockCC.templateWriter = &writers.LocalFSWriter{}
+
+	detectedLang, lowerLang, err := mockCC.mockDetectLanguage()
+	assert.False(t, detectedLang == nil)
+	assert.True(t, lowerLang == "python")
+	assert.Nil(t, err)
+
+	err = mockCC.generateDockerfile(detectedLang, lowerLang)
+	assert.True(t, err == nil)
+
+	dockerFileContent, err := ioutil.ReadFile("Dockerfile")
+	if err != nil {
+		t.Error(err)
+	}
+	assert.Contains(t, string(dockerFileContent), "CMD [\"main.py\"]")
+
+	err = os.Remove("Dockerfile")
+	if err != nil {
+		t.Error(err)
+	}
+	os.RemoveAll(".dockerignore")
+	if err != nil {
+		t.Error(err)
+	}
+}
+
 func TestInitConfig(t *testing.T) {
 	mockCC := &createCmd{}
 	mockCC.createConfig = &CreateConfig{}

pkg/languages/defaults/python.go

@@ -2,6 +2,8 @@ package defaults
 
 import (
 	"fmt"
+	"path/filepath"
+	"regexp"
 
 	"github.com/Azure/draft/pkg/reporeader"
 )
@@ -16,8 +18,32 @@ func (p PythonExtractor) ReadDefaults(r reporeader.RepoReader) (map[string]strin
 	if err != nil {
 		return nil, fmt.Errorf("error finding python files: %v", err)
 	}
-	if len(files) > 0 {
-		extractedValues["ENTRYPOINT"] = files[0]
+
+	entryPointPattern := `if\s*__name__\s*==\s*["']__main__["']`
+	compiledPattern, err := regexp.Compile(entryPointPattern)
+	if err != nil {
+		return nil, fmt.Errorf("error compiling regex pattern: %v", err)
+	}
+
+	for _, filePath := range files {
+		fileContent, err := r.ReadFile(filePath)
+		baseFile := filepath.Base(filePath)
+
+		if err != nil {
+			return nil, fmt.Errorf(("error reading python files"))
+		}
+		fileContentInString := string(fileContent)
+
+		if compiledPattern.MatchString(fileContentInString) || baseFile == "main.py" || baseFile == "app.py" {
+			extractedValues["ENTRYPOINT"] = baseFile
+			break
+		}
+	}
+
+	if _, ok := extractedValues["ENTRYPOINT"]; !ok {
+		if len(files) > 0 {
+			extractedValues["ENTRYPOINT"] = files[0]
+		}
 	}
 
 	return extractedValues, nil

pkg/languages/defaults/python_test.go

@@ -65,7 +65,38 @@ func TestPythonExtractor_ReadDefaults(t *testing.T) {
 				"ENTRYPOINT": "foo.py",
 			},
 			wantErr: false,
-		}, {
+		},
+		{
+			name: "extract python file containing the string \"__name__=='__main__'\" as the entrypoint",
+			args: args{
+				r: reporeader.TestRepoReader{
+					Files: map[string][]byte{
+						"foo.py": []byte("print('hello world')"),
+						"bar.py": []byte("if __name__ == '__main__': print('hello world')"),
+					},
+				},
+			},
+			want: map[string]string{
+				"ENTRYPOINT": "bar.py",
+			},
+			wantErr: false,
+		},
+		{
+			name: "extract python file named app.py as the entrypoint",
+			args: args{
+				r: reporeader.TestRepoReader{
+					Files: map[string][]byte{
+						"foo.py": []byte("print('Hello World')"),
+						"app.py": []byte("print('Hello World')"),
+					},
+				},
+			},
+			want: map[string]string{
+				"ENTRYPOINT": "app.py",
+			},
+			wantErr: false,
+		},
+		{
 			name: "no extraction if no python files",
 			args: args{
 				r: reporeader.TestRepoReader{

template/dockerfiles/python/Dockerfile

@@ -9,4 +9,4 @@ RUN pip install --no-cache-dir -r requirements.txt
 COPY . .
 
 ENTRYPOINT ["python"]
-CMD ["app.py"]
+CMD ["{{ENTRYPOINT}}"]