add test for timeout

acceptance/telemetry/upload-succeeds/script

@@ -0,0 +1,12 @@
+export DATABRICKS_CLI_TELEMETRY_PID_FILE=./telemetry.pid
+export DATABRICKS_CLI_TELEMETRY_UPLOAD_LOGS_FILE=./out.upload_process.txt
+
+trace $CLI selftest send-telemetry
+
+echo "waiting for telemetry process to finish"
+
+# Wait for the child telemetry process to finish
+wait_pid $(cat ./telemetry.pid)
+
+# cleanup the pid file
+rm -f ./telemetry.pid

acceptance/telemetry/upload-timeout/out.upload_process.txt

@@ -0,0 +1 @@
+error: Failed to flush telemetry log due to timeout

acceptance/telemetry/upload-timeout/output.txt

@@ -0,0 +1,4 @@
+
+>>> [CLI] selftest send-telemetry
+waiting for telemetry process to finish
+[wait_pid] process has ended

acceptance/telemetry/upload-timeout/script

@@ -0,0 +1,16 @@
+export DATABRICKS_CLI_TELEMETRY_PID_FILE=./telemetry.pid
+export DATABRICKS_CLI_TELEMETRY_UPLOAD_LOGS_FILE=./out.upload_process.txt
+
+# Configure a timeout of 0 seconds. This ensures that the timeout is respected and the
+# telemetry process does not try to upload logs.
+export DATABRICKS_CLI_TELEMETRY_UPLOAD_TIMEOUT="0s"
+
+trace $CLI selftest send-telemetry
+
+echo "waiting for telemetry process to finish"
+
+# Wait for the child telemetry process to finish
+wait_pid $(cat ./telemetry.pid)
+
+# cleanup the pid file
+rm -f ./telemetry.pid

acceptance/telemetry/upload-timeout/test.toml

@@ -0,0 +1,11 @@
+[[Server]]
+Pattern = "POST /telemetry-ext"
+Response.Body = '''
+{
+    "numProtoSuccess": 2
+}
+'''
+
+[[Repls]]
+Old = 'execution_time_ms\\\":\d{1,5},'
+New = 'execution_time_ms\":\"SMALL_INT\",'

cmd/root/auth.go

@@ -186,8 +186,7 @@ func workspaceClientOrPrompt(ctx context.Context, cfg *config.Config, allowPromp
 	return w, err
 }
 
-// TODO: Make upload with oauth work.
-// TODO: Move env var inheritance to the daemon library.
+// TODO: Run as integration tests?
 func MustWorkspaceClient(cmd *cobra.Command, args []string) error {
 	cfg := &config.Config{}
 

cmd/root/root.go

@@ -103,6 +103,7 @@ func flagErrorFunc(c *cobra.Command, err error) error {
 // Execute adds all child commands to the root command and sets flags appropriately.
 // This is called by main.main(). It only needs to happen once to the rootCmd.
 func Execute(ctx context.Context, cmd *cobra.Command) error {
+	// TODO: deferred panic recovery
 	ctx = telemetry.WithNewLogger(ctx)
 	ctx = dbr.DetectRuntime(ctx)
 	start := time.Now()

libs/telemetry/upload.go

@@ -25,6 +25,9 @@ const (
 	// Environment variable to disable telemetry. If this is set to any value, telemetry
 	// will be disabled.
 	DisableEnvVar = "DATABRICKS_CLI_DISABLE_TELEMETRY"
+
+	// Max time to try and upload the telemetry logs. Useful for testing.
+	UploadTimeout = "DATABRICKS_CLI_TELEMETRY_UPLOAD_TIMEOUT"
 )
 
 type UploadConfig struct {
@@ -68,8 +71,16 @@ func Upload() (*ResponseBody, error) {
 		return nil, fmt.Errorf("Failed to create API client: %s\n", err)
 	}
 
+	maxUploadTime := 30 * time.Second
+	if v, ok := os.LookupEnv(UploadTimeout); ok {
+		maxUploadTime, err = time.ParseDuration(v)
+		if err != nil {
+			return nil, fmt.Errorf("Failed to parse time limit %s: %s\n", UploadTimeout, err)
+		}
+	}
+
 	// Set a maximum total time to try telemetry uploads.
-	ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
+	ctx, cancel := context.WithTimeout(context.Background(), maxUploadTime)
 	defer cancel()
 
 	resp := &ResponseBody{}