ref: add integration test for postChunk endpoint

cmd/vroom/chunk.go

@@ -29,7 +29,9 @@ func (env *environment) postChunk(w http.ResponseWriter, r *http.Request) {
 	body, err := io.ReadAll(r.Body)
 	s.Finish()
 	if err != nil {
-		hub.CaptureException(err)
+		if hub != nil {
+			hub.CaptureException(err)
+		}
 		w.WriteHeader(http.StatusBadRequest)
 		return
 	}
@@ -41,34 +43,40 @@ func (env *environment) postChunk(w http.ResponseWriter, r *http.Request) {
 	err = json.Unmarshal(body, c)
 	s.Finish()
 	if err != nil {
-		hub.CaptureException(err)
+		if hub != nil {
+			hub.CaptureException(err)
+		}
 		w.WriteHeader(http.StatusBadRequest)
 		return
 	}
 
-	hub.Scope().SetContext("Profile metadata", map[string]interface{}{
-		"chunk_id":        c.ID,
-		"organization_id": strconv.FormatUint(c.OrganizationID, 10),
-		"profiler_id":     c.ProfilerID,
-		"project_id":      strconv.FormatUint(c.ProjectID, 10),
-		"size":            len(body),
-	})
+	if hub != nil {
+		hub.Scope().SetContext("Profile metadata", map[string]interface{}{
+			"chunk_id":        c.ID,
+			"organization_id": strconv.FormatUint(c.OrganizationID, 10),
+			"profiler_id":     c.ProfilerID,
+			"project_id":      strconv.FormatUint(c.ProjectID, 10),
+			"size":            len(body),
+		})
 
-	hub.Scope().SetTags(map[string]string{
-		"platform": string(c.Platform),
-	})
+		hub.Scope().SetTags(map[string]string{
+			"platform": string(c.Platform),
+		})
+	}
 
 	s = sentry.StartSpan(ctx, "gcs.write")
 	s.Description = "Write profile to GCS"
-	err = storageutil.CompressedWrite(ctx, env.storage, c.StoragePath(), body)
+	err = storageutil.CompressedWrite(ctx, env.storage, c.StoragePath(), c)
 	s.Finish()
 	if err != nil {
 		if errors.Is(err, context.DeadlineExceeded) {
 			// This is a transient error, we'll retry
 			w.WriteHeader(http.StatusTooManyRequests)
 		} else {
 			// These errors won't be retried
-			hub.CaptureException(err)
+			if hub != nil {
+				hub.CaptureException(err)
+			}
 			if code := gcerrors.Code(err); code == gcerrors.FailedPrecondition {
 				w.WriteHeader(http.StatusPreconditionFailed)
 			} else {
@@ -83,7 +91,9 @@ func (env *environment) postChunk(w http.ResponseWriter, r *http.Request) {
 	b, err := json.Marshal(buildChunkKafkaMessage(c))
 	s.Finish()
 	if err != nil {
-		hub.CaptureException(err)
+		if hub != nil {
+			hub.CaptureException(err)
+		}
 		w.WriteHeader(http.StatusInternalServerError)
 		return
 	}
@@ -94,7 +104,9 @@ func (env *environment) postChunk(w http.ResponseWriter, r *http.Request) {
 		Value: b,
 	})
 	if err != nil {
-		hub.CaptureException(err)
+		if hub != nil {
+			hub.CaptureException(err)
+		}
 		w.WriteHeader(http.StatusInternalServerError)
 		return
 	}

cmd/vroom/chunk_test.go

@@ -0,0 +1,148 @@
+package main
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"fmt"
+	"log"
+	"net/http/httptest"
+	"os"
+	"testing"
+	"time"
+
+	"github.com/getsentry/vroom/internal/chunk"
+	"github.com/getsentry/vroom/internal/frame"
+	"github.com/getsentry/vroom/internal/storageutil"
+	"github.com/getsentry/vroom/internal/testutil"
+	"github.com/google/uuid"
+	"github.com/segmentio/kafka-go"
+	"gocloud.dev/blob"
+	_ "gocloud.dev/blob/fileblob"
+	_ "gocloud.dev/blob/gcsblob"
+)
+
+var fileBlobBucket *blob.Bucket
+
+func TestMain(m *testing.M) {
+	temporaryDirectory, err := os.MkdirTemp(os.TempDir(), "sentry-profiles-*")
+	if err != nil {
+		log.Fatalf("couldn't create a temporary directory: %s", err.Error())
+	}
+
+	fileBlobBucket, err = blob.OpenBucket(context.Background(), "file://localhost/"+temporaryDirectory)
+	if err != nil {
+		log.Fatalf("couldn't open a local filesystem bucket: %s", err.Error())
+	}
+
+	code := m.Run()
+
+	if err := fileBlobBucket.Close(); err != nil {
+		log.Printf("couldn't close the local filesystem bucket: %s", err.Error())
+	}
+
+	err = os.RemoveAll(temporaryDirectory)
+	if err != nil {
+		log.Printf("couldn't remove the temporary directory: %s", err.Error())
+	}
+
+	os.Exit(code)
+}
+
+func TestPostAndReadChunk(t *testing.T) {
+	profilerID := uuid.New().String()
+	chunkID := uuid.New().String()
+	chunkData := chunk.Chunk{
+		ID:             chunkID,
+		ProfilerID:     profilerID,
+		Environment:    "dev",
+		Platform:       "python",
+		Release:        "1.2",
+		OrganizationID: 1,
+		ProjectID:      1,
+		Profile: chunk.Data{
+			Frames: []frame.Frame{
+				{Function: "test"},
+			},
+			Stacks: [][]int{
+				{0},
+			},
+			Samples: []chunk.Sample{
+				{StackID: 0, Timestamp: 1.0},
+			},
+		},
+		Measurements: json.RawMessage("null"),
+	}
+
+	objectName := fmt.Sprintf(
+		"%d/%d/%s/%s",
+		chunkData.OrganizationID,
+		chunkData.ProjectID,
+		chunkData.ProfilerID,
+		chunkData.ID,
+	)
+
+	tests := []struct {
+		name       string
+		blobBucket *blob.Bucket
+		objectName string
+	}{
+		{
+			name:       "Filesystem",
+			blobBucket: fileBlobBucket,
+			objectName: objectName,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			env := environment{
+				storage: test.blobBucket,
+				profilingWriter: &kafka.Writer{
+					Addr:         kafka.TCP("localhost:9092"),
+					Async:        true,
+					Balancer:     kafka.CRC32Balancer{},
+					BatchBytes:   20 * MiB,
+					BatchSize:    10,
+					Compression:  kafka.Lz4,
+					ReadTimeout:  3 * time.Second,
+					WriteTimeout: 3 * time.Second,
+				},
+				config: ServiceConfig{
+					ProfileChunksKafkaTopic: "snuba-profile-chunks",
+				},
+			}
+			jsonValue, err := json.Marshal(chunkData)
+			if err != nil {
+				t.Fatal(err)
+			}
+
+			req := httptest.NewRequest("POST", "/", bytes.NewBuffer(jsonValue))
+			w := httptest.NewRecorder()
+
+			// POST the chunk and check the we get a 204 response status code
+			env.postChunk(w, req)
+			resp := w.Result()
+			defer resp.Body.Close()
+			if resp.StatusCode != 204 {
+				t.Fatalf("Expected status code 204. Found: %d", resp.StatusCode)
+			}
+
+			// read the chunk with UnmarshalCompressed and make check that we can unmarshal
+			// the data into the Chunk struct and that it matches the original
+			var c chunk.Chunk
+			err = storageutil.UnmarshalCompressed(
+				context.Background(),
+				test.blobBucket,
+				objectName,
+				&c,
+			)
+			if err != nil {
+				t.Fatal(err)
+			}
+			if diff := testutil.Diff(chunkData, c); diff != "" {
+				t.Fatalf("Result mismatch: got - want +\n%s", diff)
+			}
+		})
+	}
+}