Add support for /workspace/import in testserver (#3192)

## Changes
1. Adds /worksapce/import to the default test server.
2. Add support for Python source notebooks to the test server.

## Why
It's useful in converting generate unit tests to acceptance tests. PR:
#3188

## Tests
Existing tests pass.

Author: shreyas-goenka

acceptance/internal/handlers.go

@@ -1,6 +1,7 @@
 package internal
 
 import (
+	"encoding/base64"
 	"encoding/json"
 	"fmt"
 	"net/http"
@@ -113,6 +114,37 @@ func addDefaultHandlers(server *testserver.Server) {
 		return ""
 	})
 
+	server.Handle("POST", "/api/2.0/workspace/import", func(req testserver.Request) any {
+		var request workspace.Import
+		err := json.Unmarshal(req.Body, &request)
+		if err != nil {
+			return testserver.Response{
+				Body:       fmt.Sprintf("internal error: %s", err),
+				StatusCode: http.StatusInternalServerError,
+			}
+		}
+
+		if request.Format != workspace.ImportFormatAuto {
+			return testserver.Response{
+				Body:       "internal error: The test server only supports auto format.",
+				StatusCode: http.StatusInternalServerError,
+			}
+		}
+
+		// The /workspace/import endpoint expects the content as base64 encoded string.
+		// We need to decode it to get the actual content.
+		decoded, err := base64.StdEncoding.DecodeString(request.Content)
+		if err != nil {
+			return testserver.Response{
+				Body:       fmt.Sprintf("internal error: %s", err),
+				StatusCode: http.StatusInternalServerError,
+			}
+		}
+
+		req.Workspace.WorkspaceFilesImportFile(request.Path, decoded)
+		return ""
+	})
+
 	server.Handle("GET", "/api/2.0/workspace-files/{path:.*}", func(req testserver.Request) any {
 		path := req.Vars["path"]
 		return req.Workspace.WorkspaceFilesExportFile(path)

libs/testserver/fake_workspace.go

@@ -5,6 +5,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"path"
+	"path/filepath"
 	"sort"
 	"strconv"
 	"strings"
@@ -179,17 +180,38 @@ func (s *FakeWorkspace) WorkspaceFilesImportFile(filePath string, body []byte) {
 
 	defer s.LockUnlock()()
 
-	s.files[filePath] = FileEntry{
-		Info: workspace.ObjectInfo{
-			ObjectType: "FILE",
-			Path:       filePath,
-			Language:   "SCALA",
-		},
-		Data: body,
+	workspacePath := filePath
+
+	// Note: Files with .py, .scala, .r or .sql extension can
+	// be notebooks if they contain a magical "Databricks notebook source"
+	// header comment. We omit support non-python extensions for now for simplicity.
+	extension := filepath.Ext(filePath)
+	if extension == ".py" && strings.HasPrefix(string(body), "# Databricks notebook source") {
+		// Notebooks are stripped of their extension by the workspace import API.
+		workspacePath = strings.TrimSuffix(filePath, extension)
+		s.files[workspacePath] = FileEntry{
+			Info: workspace.ObjectInfo{
+				ObjectType: "NOTEBOOK",
+				Path:       workspacePath,
+				Language:   "PYTHON",
+			},
+			Data: body,
+		}
+	} else {
+		// The endpoint does not set language for files, so we omit that
+		// here as well.
+		// ref: https://docs.databricks.com/api/workspace/workspace/getstatus#language
+		s.files[workspacePath] = FileEntry{
+			Info: workspace.ObjectInfo{
+				ObjectType: "FILE",
+				Path:       workspacePath,
+			},
+			Data: body,
+		}
 	}
 
 	// Add all directories in the path to the directories map
-	for dir := path.Dir(filePath); dir != "/"; dir = path.Dir(dir) {
+	for dir := path.Dir(workspacePath); dir != "/"; dir = path.Dir(dir) {
 		s.directories[dir] = true
 	}
 }