This commit is contained in:
Shreyas Goenka 2025-02-18 16:35:14 +01:00
parent 520f06905b
commit a98993abfc
No known key found for this signature in database
GPG Key ID: 92A07DF49CCB0622
1 changed files with 8 additions and 9 deletions

View File

@ -2,7 +2,6 @@ package telemetry
import ( import (
"encoding/json" "encoding/json"
"net/http"
"os" "os"
"path/filepath" "path/filepath"
"testing" "testing"
@ -15,24 +14,24 @@ import (
"github.com/stretchr/testify/require" "github.com/stretchr/testify/require"
) )
func TestTelemetryUpload(t *testing.T) { func TestTelemetryUploadRetries(t *testing.T) {
server := testserver.New(t) server := testserver.New(t)
t.Cleanup(server.Close) t.Cleanup(server.Close)
count := 0 count := 0
server.Handle("POST", "/telemetry-ext", func(_ *testserver.FakeWorkspace, req *http.Request) (resp any, statusCode int) { server.Handle("POST", "/telemetry-ext", func(req testserver.Request) any {
count++ count++
if count == 1 { if count == 1 {
return ResponseBody{ return ResponseBody{
NumProtoSuccess: 1, NumProtoSuccess: 1,
}, http.StatusOK }
} }
if count == 2 { if count == 2 {
return ResponseBody{ return ResponseBody{
NumProtoSuccess: 2, NumProtoSuccess: 2,
}, http.StatusOK }
} }
return nil, http.StatusInternalServerError return nil
}) })
t.Setenv("DATABRICKS_HOST", server.URL) t.Setenv("DATABRICKS_HOST", server.URL)
@ -67,14 +66,14 @@ func TestTelemetryUpload(t *testing.T) {
tmpDir := t.TempDir() tmpDir := t.TempDir()
testutil.WriteFile(t, filepath.Join(tmpDir, "stdin"), string(b)) testutil.WriteFile(t, filepath.Join(tmpDir, "stdin"), string(b))
fd, err := os.OpenFile(filepath.Join(tmpDir, "stdin"), os.O_RDONLY, 0o644) f, err := os.OpenFile(filepath.Join(tmpDir, "stdin"), os.O_RDONLY, 0o644)
require.NoError(t, err) require.NoError(t, err)
// Redirect stdin to the file containing the telemetry logs. // Redirect stdin to the file containing the telemetry logs.
old := os.Stdin old := os.Stdin
os.Stdin = fd os.Stdin = f
t.Cleanup(func() { t.Cleanup(func() {
fd.Close() f.Close()
os.Stdin = old os.Stdin = old
}) })