fixup test

Emyrk · Emyrk · commit 4ec22d8d0f96 · 2025-06-09T15:36:05.000-05:00
diff --git a/coderd/provisionerdserver/upload_file_test.go b/coderd/provisionerdserver/upload_file_test.go
@@ -0,0 +1,196 @@
+package provisionerdserver_test
+
+import (
+	"context"
+	crand "crypto/rand"
+	"fmt"
+	"testing"
+
+	"github.com/google/uuid"
+	"github.com/stretchr/testify/require"
+	"golang.org/x/xerrors"
+	"storj.io/drpc"
+
+	"github.com/coder/coder/v2/coderd/database"
+	"github.com/coder/coder/v2/coderd/database/dbmem"
+	"github.com/coder/coder/v2/coderd/externalauth"
+	"github.com/coder/coder/v2/codersdk/drpcsdk"
+	proto "github.com/coder/coder/v2/provisionerd/proto"
+	sdkproto "github.com/coder/coder/v2/provisionersdk/proto"
+	"github.com/coder/coder/v2/testutil"
+)
+
+// TestUploadFileLargeModuleFiles tests the UploadFile RPC with large module files
+func TestUploadFileLargeModuleFiles(t *testing.T) {
+	t.Parallel()
+
+	ctx := testutil.Context(t, testutil.WaitMedium)
+	db := dbmem.New()
+
+	// Create server
+	server, db, _, _ := setup(t, false, &overrides{
+		externalAuthConfigs: []*externalauth.Config{{}},
+	})
+
+	testSizes := []int{
+		512,                           // A small file
+		drpcsdk.MaxMessageSize + 1024, // Just over the limit
+		drpcsdk.MaxMessageSize * 2,    // 2x the limit
+		sdkproto.ChunkSize*3 + 512,    // Multiple chunks with partial last
+	}
+
+	for _, size := range testSizes {
+		t.Run(fmt.Sprintf("size_%d_bytes", size), func(t *testing.T) {
+			// Generate test module files data
+			moduleData := make([]byte, size)
+			_, err := crand.Read(moduleData)
+			require.NoError(t, err)
+
+			// Convert to upload format
+			upload, chunks := sdkproto.BytesToDataUpload(sdkproto.DataUploadType_UPLOAD_TYPE_MODULE_FILES, moduleData)
+
+			stream := newMockUploadStream(upload, chunks...)
+
+			// Execute upload
+			err = server.UploadFile(stream)
+			require.NoError(t, err)
+
+			// Upload should be done
+			require.True(t, stream.isDone(), "stream should be done after upload")
+
+			// Verify file was stored in database
+			hashString := fmt.Sprintf("%x", upload.DataHash)
+			file, err := db.GetFileByHashAndCreator(ctx, database.GetFileByHashAndCreatorParams{
+				Hash:      hashString,
+				CreatedBy: uuid.Nil, // Provisionerd creates with Nil UUID
+			})
+			require.NoError(t, err)
+			require.Equal(t, hashString, file.Hash)
+			require.Equal(t, moduleData, file.Data)
+			require.Equal(t, "application/x-tar", file.Mimetype)
+
+			// Try to upload it again, and it should still be successful
+			stream = newMockUploadStream(upload, chunks...)
+			err = server.UploadFile(stream)
+			require.NoError(t, err, "re-upload should succeed without error")
+			require.True(t, stream.isDone(), "stream should be done after re-upload")
+		})
+	}
+}
+
+// TestUploadFileErrorScenarios tests various error conditions in file upload
+func TestUploadFileErrorScenarios(t *testing.T) {
+	t.Parallel()
+
+	server, _, _, _ := setup(t, false, &overrides{
+		externalAuthConfigs: []*externalauth.Config{{}},
+	})
+
+	// Generate test data
+	moduleData := make([]byte, sdkproto.ChunkSize*2)
+	_, err := crand.Read(moduleData)
+	require.NoError(t, err)
+
+	upload, chunks := sdkproto.BytesToDataUpload(sdkproto.DataUploadType_UPLOAD_TYPE_MODULE_FILES, moduleData)
+	var _ = upload
+
+	t.Run("chunk_before_upload", func(t *testing.T) {
+		stream := newMockUploadStream(nil, chunks[0])
+
+		err := server.UploadFile(stream)
+		require.ErrorContains(t, err, "unexpected chunk piece while waiting for file upload")
+		require.True(t, stream.isDone(), "stream should be done after error")
+	})
+
+	t.Run("duplicate_upload", func(t *testing.T) {
+		stream := &mockUploadStream{
+			done:     make(chan struct{}),
+			messages: make(chan *proto.UploadFileRequest, 2),
+		}
+
+		up := &proto.UploadFileRequest{Type: &proto.UploadFileRequest_DataUpload{&proto.DataUpload{
+			UploadType: proto.DataUploadType(upload.UploadType),
+			DataHash:   upload.DataHash,
+			FileSize:   upload.FileSize,
+			Chunks:     upload.Chunks,
+		}}}
+
+		// Send it twice
+		stream.messages <- up
+		stream.messages <- up
+
+		err := server.UploadFile(stream)
+		require.ErrorContains(t, err, "unexpected file upload while waiting for file completion")
+		require.True(t, stream.isDone(), "stream should be done after error")
+	})
+
+	t.Run("unsupported_upload_type", func(t *testing.T) {
+		cpy := *upload
+		cpy.UploadType = sdkproto.DataUploadType_UPLOAD_TYPE_UNKNOWN // Set to an unsupported type
+		stream := newMockUploadStream(&cpy, chunks...)
+
+		err := server.UploadFile(stream)
+		require.ErrorContains(t, err, "unsupported file upload type")
+		require.True(t, stream.isDone(), "stream should be done after error")
+	})
+}
+
+type mockUploadStream struct {
+	done     chan struct{}
+	messages chan *proto.UploadFileRequest
+}
+
+func (m mockUploadStream) SendAndClose(empty *proto.Empty) error {
+	close(m.done)
+	return nil
+}
+
+func (m mockUploadStream) Recv() (*proto.UploadFileRequest, error) {
+	msg, ok := <-m.messages
+	if !ok {
+		return nil, xerrors.New("no more messages to receive")
+	}
+	return msg, nil
+}
+func (m *mockUploadStream) Context() context.Context { panic(errUnimplemented) }
+func (m *mockUploadStream) MsgSend(msg drpc.Message, enc drpc.Encoding) error {
+	panic(errUnimplemented)
+}
+func (m *mockUploadStream) MsgRecv(msg drpc.Message, enc drpc.Encoding) error {
+	panic(errUnimplemented)
+}
+func (m *mockUploadStream) CloseSend() error { panic(errUnimplemented) }
+func (m *mockUploadStream) Close() error     { panic(errUnimplemented) }
+func (m *mockUploadStream) isDone() bool {
+	select {
+	case <-m.done:
+		return true
+	default:
+		return false
+	}
+}
+
+func newMockUploadStream(up *sdkproto.DataUpload, chunks ...*sdkproto.ChunkPiece) *mockUploadStream {
+	stream := &mockUploadStream{
+		done:     make(chan struct{}),
+		messages: make(chan *proto.UploadFileRequest, 1+len(chunks)),
+	}
+	if up != nil {
+		stream.messages <- &proto.UploadFileRequest{Type: &proto.UploadFileRequest_DataUpload{&proto.DataUpload{
+			UploadType: proto.DataUploadType(up.UploadType),
+			DataHash:   up.DataHash,
+			FileSize:   up.FileSize,
+			Chunks:     up.Chunks,
+		}}}
+	}
+
+	for _, chunk := range chunks {
+		stream.messages <- &proto.UploadFileRequest{Type: &proto.UploadFileRequest_ChunkPiece{&proto.ChunkPiece{
+			Data:         chunk.Data,
+			FullDataHash: chunk.FullDataHash,
+			PieceIndex:   chunk.PieceIndex,
+		}}}
+	}
+	close(stream.messages)
+	return stream
+}
