Merged in feature/track-filesize (pull request #200)

track file sizes for all documents in system

* feature complete

needs dev testing

internal/database/migrations/00000000000118_add_document_file_size.down.sql

@@ -0,0 +1,2 @@
+-- Remove file_size_bytes column from documents table
+ALTER TABLE documents DROP COLUMN file_size_bytes;

internal/database/migrations/00000000000118_add_document_file_size.up.sql

@@ -0,0 +1,4 @@
+-- Add file_size_bytes column to documents table
+-- Stores the size of the document in bytes as measured from S3 HeadObject
+-- NULL for legacy documents (before this feature was added)
+ALTER TABLE documents ADD COLUMN file_size_bytes BIGINT;

internal/database/queries/document.sql

@@ -34,7 +34,7 @@ SELECT id, hash from documents where clientId = @clientId;
 SELECT id, hash, filename from documents where batch_id = @batch_id;
 
 -- name: CreateDocument :one
-INSERT INTO documents (clientId, hash, batch_id, filename, folderId, originalPath) VALUES ($1, $2, $3, $4, $5, $6) RETURNING id;
+INSERT INTO documents (clientId, hash, batch_id, filename, folderId, originalPath, file_size_bytes) VALUES ($1, $2, $3, $4, $5, $6, $7) RETURNING id;
 
 -- name: AddDocumentEntry :exec
 INSERT INTO documentEntries (documentId, bucket, key) VALUES ($1, $2, $3);
@@ -85,6 +85,7 @@ SELECT
     d.hash,
     d.folderId,
     d.filename,
-    d.originalPath
+    d.originalPath,
+    d.file_size_bytes
 FROM documents d
 WHERE d.id = $1;

internal/database/queries/folders.sql

@@ -95,7 +95,8 @@ SELECT
     d.hash,
     d.folderId,
     d.filename,
-    d.originalPath
+    d.originalPath,
+    d.file_size_bytes
 FROM documents d
 WHERE d.folderId = $1
 ORDER BY d.id;

internal/database/repository/document.sql.go

@@ -65,21 +65,22 @@ func (q *Queries) AddDocumentUpload(ctx context.Context, arg *AddDocumentUploadP
 }
 
 const createDocument = `-- name: CreateDocument :one
-INSERT INTO documents (clientId, hash, batch_id, filename, folderId, originalPath) VALUES ($1, $2, $3, $4, $5, $6) RETURNING id
+INSERT INTO documents (clientId, hash, batch_id, filename, folderId, originalPath, file_size_bytes) VALUES ($1, $2, $3, $4, $5, $6, $7) RETURNING id
 `
 
 type CreateDocumentParams struct {
-	Clientid     string     `db:"clientid"`
-	Hash         string     `db:"hash"`
-	BatchID      *uuid.UUID `db:"batch_id"`
-	Filename     *string    `db:"filename"`
-	Folderid     *uuid.UUID `db:"folderid"`
-	Originalpath *string    `db:"originalpath"`
+	Clientid      string     `db:"clientid"`
+	Hash          string     `db:"hash"`
+	BatchID       *uuid.UUID `db:"batch_id"`
+	Filename      *string    `db:"filename"`
+	Folderid      *uuid.UUID `db:"folderid"`
+	Originalpath  *string    `db:"originalpath"`
+	FileSizeBytes *int64     `db:"file_size_bytes"`
 }
 
 // CreateDocument
 //
-//	INSERT INTO documents (clientId, hash, batch_id, filename, folderId, originalPath) VALUES ($1, $2, $3, $4, $5, $6) RETURNING id
+//	INSERT INTO documents (clientId, hash, batch_id, filename, folderId, originalPath, file_size_bytes) VALUES ($1, $2, $3, $4, $5, $6, $7) RETURNING id
 func (q *Queries) CreateDocument(ctx context.Context, arg *CreateDocumentParams) (uuid.UUID, error) {
 	row := q.db.QueryRow(ctx, createDocument,
 		arg.Clientid,
@@ -88,6 +89,7 @@ func (q *Queries) CreateDocument(ctx context.Context, arg *CreateDocumentParams)
 		arg.Filename,
 		arg.Folderid,
 		arg.Originalpath,
+		arg.FileSizeBytes,
 	)
 	var id uuid.UUID
 	err := row.Scan(&id)
@@ -101,18 +103,20 @@ SELECT
     d.hash,
     d.folderId,
     d.filename,
-    d.originalPath
+    d.originalPath,
+    d.file_size_bytes
 FROM documents d
 WHERE d.id = $1
 `
 
 type GetDocumentEnrichedRow struct {
-	ID           uuid.UUID  `db:"id"`
-	Clientid     string     `db:"clientid"`
-	Hash         string     `db:"hash"`
-	Folderid     *uuid.UUID `db:"folderid"`
-	Filename     *string    `db:"filename"`
-	Originalpath *string    `db:"originalpath"`
+	ID            uuid.UUID  `db:"id"`
+	Clientid      string     `db:"clientid"`
+	Hash          string     `db:"hash"`
+	Folderid      *uuid.UUID `db:"folderid"`
+	Filename      *string    `db:"filename"`
+	Originalpath  *string    `db:"originalpath"`
+	FileSizeBytes *int64     `db:"file_size_bytes"`
 }
 
 // Retrieves a document with extended metadata including folder and filename
@@ -123,7 +127,8 @@ type GetDocumentEnrichedRow struct {
 //	    d.hash,
 //	    d.folderId,
 //	    d.filename,
-//	    d.originalPath
+//	    d.originalPath,
+//	    d.file_size_bytes
 //	FROM documents d
 //	WHERE d.id = $1
 func (q *Queries) GetDocumentEnriched(ctx context.Context, id uuid.UUID) (*GetDocumentEnrichedRow, error) {
@@ -136,6 +141,7 @@ func (q *Queries) GetDocumentEnriched(ctx context.Context, id uuid.UUID) (*GetDo
 		&i.Folderid,
 		&i.Filename,
 		&i.Originalpath,
+		&i.FileSizeBytes,
 	)
 	return &i, err
 }

internal/database/repository/folders.sql.go

@@ -89,14 +89,14 @@ func (q *Queries) GetClientRootFolder(ctx context.Context, clientid string) (*Fo
 }
 
 const getDocumentsByFolder = `-- name: GetDocumentsByFolder :many
-SELECT id, clientid, hash, batch_id, filename, folderid, originalpath FROM documents
+SELECT id, clientid, hash, batch_id, filename, folderid, originalpath, file_size_bytes FROM documents
 WHERE folderId = $1
 ORDER BY id
 `
 
 // GetDocumentsByFolder
 //
-//	SELECT id, clientid, hash, batch_id, filename, folderid, originalpath FROM documents
+//	SELECT id, clientid, hash, batch_id, filename, folderid, originalpath, file_size_bytes FROM documents
 //	WHERE folderId = $1
 //	ORDER BY id
 func (q *Queries) GetDocumentsByFolder(ctx context.Context, folderid *uuid.UUID) ([]*Document, error) {
@@ -116,6 +116,7 @@ func (q *Queries) GetDocumentsByFolder(ctx context.Context, folderid *uuid.UUID)
 			&i.Filename,
 			&i.Folderid,
 			&i.Originalpath,
+			&i.FileSizeBytes,
 		); err != nil {
 			return nil, err
 		}
@@ -133,18 +134,20 @@ SELECT
     d.hash,
     d.folderId,
     d.filename,
-    d.originalPath
+    d.originalPath,
+    d.file_size_bytes
 FROM documents d
 WHERE d.folderId = $1
 ORDER BY d.id
 `
 
 type GetDocumentsByFolderEnrichedRow struct {
-	ID           uuid.UUID  `db:"id"`
-	Hash         string     `db:"hash"`
-	Folderid     *uuid.UUID `db:"folderid"`
-	Filename     *string    `db:"filename"`
-	Originalpath *string    `db:"originalpath"`
+	ID            uuid.UUID  `db:"id"`
+	Hash          string     `db:"hash"`
+	Folderid      *uuid.UUID `db:"folderid"`
+	Filename      *string    `db:"filename"`
+	Originalpath  *string    `db:"originalpath"`
+	FileSizeBytes *int64     `db:"file_size_bytes"`
 }
 
 // Gets documents in a folder with all enriched metadata fields
@@ -154,7 +157,8 @@ type GetDocumentsByFolderEnrichedRow struct {
 //	    d.hash,
 //	    d.folderId,
 //	    d.filename,
-//	    d.originalPath
+//	    d.originalPath,
+//	    d.file_size_bytes
 //	FROM documents d
 //	WHERE d.folderId = $1
 //	ORDER BY d.id
@@ -173,6 +177,7 @@ func (q *Queries) GetDocumentsByFolderEnriched(ctx context.Context, folderid *uu
 			&i.Folderid,
 			&i.Filename,
 			&i.Originalpath,
+			&i.FileSizeBytes,
 		); err != nil {
 			return nil, err
 		}

internal/database/repository/labels.sql.go

@@ -191,7 +191,7 @@ func (q *Queries) GetDocumentLabels(ctx context.Context, documentid uuid.UUID) (
 }
 
 const getDocumentsByLabel = `-- name: GetDocumentsByLabel :many
-SELECT DISTINCT d.id, d.clientid, d.hash, d.batch_id, d.filename, d.folderid, d.originalpath
+SELECT DISTINCT d.id, d.clientid, d.hash, d.batch_id, d.filename, d.folderid, d.originalpath, d.file_size_bytes
 FROM documents d
 INNER JOIN documentLabels dl ON d.id = dl.documentId
 WHERE d.clientId = $1 AND dl.label = $2
@@ -205,7 +205,7 @@ type GetDocumentsByLabelParams struct {
 
 // GetDocumentsByLabel
 //
-//	SELECT DISTINCT d.id, d.clientid, d.hash, d.batch_id, d.filename, d.folderid, d.originalpath
+//	SELECT DISTINCT d.id, d.clientid, d.hash, d.batch_id, d.filename, d.folderid, d.originalpath, d.file_size_bytes
 //	FROM documents d
 //	INNER JOIN documentLabels dl ON d.id = dl.documentId
 //	WHERE d.clientId = $1 AND dl.label = $2
@@ -227,6 +227,7 @@ func (q *Queries) GetDocumentsByLabel(ctx context.Context, arg *GetDocumentsByLa
 			&i.Filename,
 			&i.Folderid,
 			&i.Originalpath,
+			&i.FileSizeBytes,
 		); err != nil {
 			return nil, err
 		}
@@ -239,7 +240,7 @@ func (q *Queries) GetDocumentsByLabel(ctx context.Context, arg *GetDocumentsByLa
 }
 
 const getDocumentsByLabelAndFolder = `-- name: GetDocumentsByLabelAndFolder :many
-SELECT DISTINCT d.id, d.clientid, d.hash, d.batch_id, d.filename, d.folderid, d.originalpath
+SELECT DISTINCT d.id, d.clientid, d.hash, d.batch_id, d.filename, d.folderid, d.originalpath, d.file_size_bytes
 FROM documents d
 INNER JOIN documentLabels dl ON d.id = dl.documentId
 WHERE d.folderId = $1 AND dl.label = $2
@@ -253,7 +254,7 @@ type GetDocumentsByLabelAndFolderParams struct {
 
 // GetDocumentsByLabelAndFolder
 //
-//	SELECT DISTINCT d.id, d.clientid, d.hash, d.batch_id, d.filename, d.folderid, d.originalpath
+//	SELECT DISTINCT d.id, d.clientid, d.hash, d.batch_id, d.filename, d.folderid, d.originalpath, d.file_size_bytes
 //	FROM documents d
 //	INNER JOIN documentLabels dl ON d.id = dl.documentId
 //	WHERE d.folderId = $1 AND dl.label = $2
@@ -275,6 +276,7 @@ func (q *Queries) GetDocumentsByLabelAndFolder(ctx context.Context, arg *GetDocu
 			&i.Filename,
 			&i.Folderid,
 			&i.Originalpath,
+			&i.FileSizeBytes,
 		); err != nil {
 			return nil, err
 		}

internal/database/repository/models.go

@@ -425,7 +425,8 @@ type Document struct {
 	Filename *string    `db:"filename"`
 	Folderid *uuid.UUID `db:"folderid"`
 	// Original path provided during upload. IMMUTABLE after creation - never modify this value.
-	Originalpath *string `db:"originalpath"`
+	Originalpath  *string `db:"originalpath"`
+	FileSizeBytes *int64  `db:"file_size_bytes"`
 }
 
 type Documentclean struct {

internal/document/get.go

@@ -103,6 +103,7 @@ func (s *Service) GetEnriched(ctx context.Context, id uuid.UUID, includeTextReco
 		Filename:      docEnriched.Filename,
 		OriginalPath:  docEnriched.Originalpath,
 		Labels:        labels,
+		FileSizeBytes: docEnriched.FileSizeBytes,
 	}
 
 	// 5. Optionally fetch full text record if requested and exists

internal/document/init/create.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"database/sql"
 	"errors"
+	"fmt"
 	"log/slog"
 
 	docsyncrunner "queryorchestration/api/docSyncRunner"
@@ -11,6 +12,8 @@ import (
 	"queryorchestration/internal/serviceconfig/objectstore"
 	"queryorchestration/internal/serviceconfig/queue"
 
+	"github.com/aws/aws-sdk-go-v2/aws"
+	"github.com/aws/aws-sdk-go-v2/service/s3"
 	"github.com/google/uuid"
 )
 
@@ -85,16 +88,24 @@ func (s *Service) getCreateParams(ctx context.Context, doc *Create) (*createDocu
 func (s *Service) submitCreate(ctx context.Context, params *createDocumentParams) (uuid.UUID, error) {
 	var id uuid.UUID
 
-	err := s.cfg.ExecuteDBTransaction(ctx, func(ctx context.Context, q *repository.Queries) error {
+	// Measure file size from S3 (trusted source - we measure it ourselves)
+	// The S3 client must be initialized - file size tracking is required
+	fileSizeBytes, err := s.measureFileSize(ctx, params.Bucket, params.Key.String())
+	if err != nil {
+		return uuid.Nil, fmt.Errorf("failed to measure file size: %w", err)
+	}
+
+	err = s.cfg.ExecuteDBTransaction(ctx, func(ctx context.Context, q *repository.Queries) error {
 		var dbid uuid.UUID
 		if params.ID == nil {
 			createid, err := s.cfg.GetDBQueries().CreateDocument(ctx, &repository.CreateDocumentParams{
-				Clientid:     params.Key.ClientID,
-				Hash:         params.Hash,
-				BatchID:      params.Key.BatchID,
-				Filename:     params.Filename,
-				Folderid:     params.FolderID,
-				Originalpath: params.OriginalPath,
+				Clientid:      params.Key.ClientID,
+				Hash:          params.Hash,
+				BatchID:       params.Key.BatchID,
+				Filename:      params.Filename,
+				Folderid:      params.FolderID,
+				Originalpath:  params.OriginalPath,
+				FileSizeBytes: fileSizeBytes,
 			})
 			if err != nil {
 				return err
@@ -127,3 +138,38 @@ func (s *Service) submitCreate(ctx context.Context, params *createDocumentParams
 
 	return id, nil
 }
+
+// measureFileSize retrieves the file size from S3 using HeadObject.
+// Returns the file size in bytes, or an error if the S3 client is not configured
+// or if the HeadObject call fails.
+//
+// Parameters:
+//   - ctx: Context for the S3 operation
+//   - bucket: S3 bucket name
+//   - key: S3 object key
+//
+// Returns:
+//   - *int64: Pointer to file size in bytes
+//   - error: Error if S3 client is nil or HeadObject fails
+//
+// Note: This operation requires s3:GetObject permission on the bucket/key.
+func (s *Service) measureFileSize(ctx context.Context, bucket, key string) (*int64, error) {
+	storeClient := s.cfg.GetStoreClient()
+	if storeClient == nil {
+		return nil, errors.New("S3 client not initialized - required for file size measurement")
+	}
+
+	headOutput, err := storeClient.HeadObject(ctx, &s3.HeadObjectInput{
+		Bucket: aws.String(bucket),
+		Key:    aws.String(key),
+	})
+	if err != nil {
+		return nil, fmt.Errorf("HeadObject failed for s3://%s/%s: %w", bucket, key, err)
+	}
+
+	if headOutput.ContentLength == nil {
+		return nil, fmt.Errorf("HeadObject returned nil ContentLength for s3://%s/%s", bucket, key)
+	}
+
+	return headOutput.ContentLength, nil
+}

internal/document/init/create_test.go

@@ -1,6 +1,7 @@
 package documentinit
 
 import (
+	"bytes"
 	"errors"
 	"fmt"
 	"testing"
@@ -9,10 +10,14 @@ import (
 	"queryorchestration/internal/database/repository"
 	"queryorchestration/internal/document"
 	"queryorchestration/internal/serviceconfig/objectstore"
+	"queryorchestration/internal/test"
+	objectstoremock "queryorchestration/mocks/objectstore"
 	queuemock "queryorchestration/mocks/queue"
 
 	"github.com/stretchr/testify/require"
 
+	"github.com/aws/aws-sdk-go-v2/aws"
+	"github.com/aws/aws-sdk-go-v2/service/s3"
 	"github.com/aws/aws-sdk-go-v2/service/sqs"
 	"github.com/google/uuid"
 	"github.com/pashagolub/pgxmock/v3"
@@ -30,6 +35,11 @@ func TestCreate(t *testing.T) {
 	cfg.DBQueries = repository.New(pool)
 	mockSQS := queuemock.NewMockSQSClient(t)
 	cfg.QueueClient = mockSQS
+
+	// Set up mock S3 client for file size measurement
+	mockS3 := objectstoremock.NewMockS3Client(t)
+	cfg.StoreClient = mockS3
+
 	cfg.DocumentSyncURL = "/i/am/here"
 
 	svc := New(cfg)
@@ -48,11 +58,19 @@ func TestCreate(t *testing.T) {
 		Location:  objectstore.Import,
 	}
 
+	// Mock HeadObject to return file size
+	fileSize := int64(1024)
+	mockS3.EXPECT().
+		HeadObject(mock.Anything, mock.MatchedBy(func(in *s3.HeadObjectInput) bool {
+			return *in.Bucket == bucket && *in.Key == key.String()
+		})).
+		Return(&s3.HeadObjectOutput{ContentLength: &fileSize}, nil)
+
 	pool.ExpectQuery("name: GetDocumentIDByHash :one").WithArgs(doc.Hash, doc.ClientID).WillReturnRows(
 		pgxmock.NewRows([]string{"id"}),
 	)
 	pool.ExpectBegin()
-	pool.ExpectQuery("name: CreateDocument :one").WithArgs(doc.ClientID, doc.Hash, (*uuid.UUID)(nil), (*string)(nil), (*uuid.UUID)(nil), (*string)(nil)).
+	pool.ExpectQuery("name: CreateDocument :one").WithArgs(doc.ClientID, doc.Hash, (*uuid.UUID)(nil), (*string)(nil), (*uuid.UUID)(nil), (*string)(nil), &fileSize).
 		WillReturnRows(
 			pgxmock.NewRows([]string{"id"}).
 				AddRow(doc.ID),
@@ -220,6 +238,10 @@ func TestSubmitCreate(t *testing.T) {
 	cfg.DBPool = pool
 	cfg.DBQueries = repository.New(pool)
 
+	// Set up mock S3 client for file size measurement
+	mockS3 := objectstoremock.NewMockS3Client(t)
+	cfg.StoreClient = mockS3
+
 	svc := New(cfg)
 
 	doc := document.DocumentSummary{
@@ -235,8 +257,16 @@ func TestSubmitCreate(t *testing.T) {
 		Location:  objectstore.Import,
 	}
 
+	// Mock HeadObject to return file size
+	fileSize := int64(2048)
+	mockS3.EXPECT().
+		HeadObject(mock.Anything, mock.MatchedBy(func(in *s3.HeadObjectInput) bool {
+			return *in.Bucket == bucket && *in.Key == key.String()
+		})).
+		Return(&s3.HeadObjectOutput{ContentLength: &fileSize}, nil)
+
 	pool.ExpectBegin()
-	pool.ExpectQuery("name: CreateDocument :one").WithArgs(doc.ClientID, doc.Hash, (*uuid.UUID)(nil), (*string)(nil), (*uuid.UUID)(nil), (*string)(nil)).
+	pool.ExpectQuery("name: CreateDocument :one").WithArgs(doc.ClientID, doc.Hash, (*uuid.UUID)(nil), (*string)(nil), (*uuid.UUID)(nil), (*string)(nil), &fileSize).
 		WillReturnRows(
 			pgxmock.NewRows([]string{"id"}).
 				AddRow(doc.ID),
@@ -264,6 +294,10 @@ func TestSubmitCreateExists(t *testing.T) {
 	cfg.DBPool = pool
 	cfg.DBQueries = repository.New(pool)
 
+	// Set up mock S3 client for file size measurement
+	mockS3 := objectstoremock.NewMockS3Client(t)
+	cfg.StoreClient = mockS3
+
 	svc := New(cfg)
 
 	doc := document.DocumentSummary{
@@ -279,6 +313,14 @@ func TestSubmitCreateExists(t *testing.T) {
 		Location:  objectstore.Import,
 	}
 
+	// Mock HeadObject to return file size (still called even for existing docs)
+	fileSize := int64(4096)
+	mockS3.EXPECT().
+		HeadObject(mock.Anything, mock.MatchedBy(func(in *s3.HeadObjectInput) bool {
+			return *in.Bucket == bucket && *in.Key == key.String()
+		})).
+		Return(&s3.HeadObjectOutput{ContentLength: &fileSize}, nil)
+
 	pool.ExpectBegin()
 	pool.ExpectExec("name: AddDocumentEntry :exec").WithArgs(doc.ID, bucket, key.String()).
 		WillReturnResult(pgxmock.NewResult("", 1))
@@ -294,3 +336,200 @@ func TestSubmitCreateExists(t *testing.T) {
 	require.NoError(t, err)
 	assert.Equal(t, doc.ID, id)
 }
+
+// TestSubmitCreate_HeadObjectError verifies that HeadObject errors are properly propagated.
+func TestSubmitCreate_HeadObjectError(t *testing.T) {
+	ctx := t.Context()
+
+	pool, err := pgxmock.NewPool()
+	require.NoError(t, err)
+	cfg := &DocInitConfig{}
+	cfg.DBPool = pool
+	cfg.DBQueries = repository.New(pool)
+
+	// Set up mock S3 client that returns an error
+	mockS3 := objectstoremock.NewMockS3Client(t)
+	cfg.StoreClient = mockS3
+
+	svc := New(cfg)
+
+	doc := document.DocumentSummary{
+		ID:       uuid.New(),
+		ClientID: "hello",
+		Hash:     "example_hash",
+	}
+	bucket := "example_bucket"
+	key := objectstore.BucketKey{
+		ClientID:  doc.ClientID,
+		CreatedAt: time.Now().UTC(),
+		EntityID:  uuid.New(),
+		Location:  objectstore.Import,
+	}
+
+	// Mock HeadObject to return an error
+	mockS3.EXPECT().
+		HeadObject(mock.Anything, mock.Anything).
+		Return(nil, errors.New("access denied"))
+
+	_, err = svc.submitCreate(ctx, &createDocumentParams{
+		Hash:   doc.Hash,
+		Bucket: bucket,
+		Key:    key,
+	})
+	assert.Error(t, err)
+	assert.Contains(t, err.Error(), "failed to measure file size")
+	assert.Contains(t, err.Error(), "access denied")
+}
+
+// TestSubmitCreate_NoS3Client verifies that missing S3 client returns an error.
+func TestSubmitCreate_NoS3Client(t *testing.T) {
+	ctx := t.Context()
+
+	pool, err := pgxmock.NewPool()
+	require.NoError(t, err)
+	cfg := &DocInitConfig{}
+	cfg.DBPool = pool
+	cfg.DBQueries = repository.New(pool)
+	// Intentionally not setting StoreClient
+
+	svc := New(cfg)
+
+	doc := document.DocumentSummary{
+		ID:       uuid.New(),
+		ClientID: "hello",
+		Hash:     "example_hash",
+	}
+	bucket := "example_bucket"
+	key := objectstore.BucketKey{
+		ClientID:  doc.ClientID,
+		CreatedAt: time.Now().UTC(),
+		EntityID:  uuid.New(),
+		Location:  objectstore.Import,
+	}
+
+	_, err = svc.submitCreate(ctx, &createDocumentParams{
+		Hash:   doc.Hash,
+		Bucket: bucket,
+		Key:    key,
+	})
+	assert.Error(t, err)
+	assert.Contains(t, err.Error(), "S3 client not initialized")
+}
+
+// TestMeasureFileSize_NilContentLength verifies handling of nil ContentLength.
+func TestMeasureFileSize_NilContentLength(t *testing.T) {
+	ctx := t.Context()
+
+	pool, err := pgxmock.NewPool()
+	require.NoError(t, err)
+	cfg := &DocInitConfig{}
+	cfg.DBPool = pool
+	cfg.DBQueries = repository.New(pool)
+
+	// Set up mock S3 client that returns nil ContentLength
+	mockS3 := objectstoremock.NewMockS3Client(t)
+	cfg.StoreClient = mockS3
+
+	svc := New(cfg)
+
+	bucket := "example_bucket"
+	key := "example_key"
+
+	mockS3.EXPECT().
+		HeadObject(mock.Anything, mock.Anything).
+		Return(&s3.HeadObjectOutput{ContentLength: nil}, nil)
+
+	_, err = svc.measureFileSize(ctx, bucket, key)
+	assert.Error(t, err)
+	assert.Contains(t, err.Error(), "nil ContentLength")
+}
+
+// TestMeasureFileSize_Integration tests the measureFileSize function using a real
+// S3 client (localstack) to verify that HeadObject correctly returns file sizes.
+// This is an integration test that requires Docker.
+func TestMeasureFileSize_Integration(t *testing.T) {
+	ctx := t.Context()
+
+	// Set up AWS container (localstack) and S3 client
+	cfg := &DocInitConfig{}
+	test.CreateAWSResources(t, cfg)
+
+	svc := New(cfg)
+
+	// Upload a test file with known content
+	bucket := cfg.GetBucket()
+	testKey := fmt.Sprintf("test/filesize/%s.txt", uuid.New().String())
+	testContent := []byte("Hello, World! This is a test file for measuring file size.")
+	expectedSize := int64(len(testContent))
+
+	// Upload the test file to S3
+	_, err := cfg.GetStoreClient().PutObject(ctx, &s3.PutObjectInput{
+		Bucket: aws.String(bucket),
+		Key:    aws.String(testKey),
+		Body:   bytes.NewReader(testContent),
+	})
+	require.NoError(t, err, "failed to upload test file to S3")
+
+	// Measure the file size using our function
+	measuredSize, err := svc.measureFileSize(ctx, bucket, testKey)
+	require.NoError(t, err, "measureFileSize should not return an error")
+	require.NotNil(t, measuredSize, "measured size should not be nil")
+
+	// Verify the measured size matches expected
+	assert.Equal(t, expectedSize, *measuredSize, "measured file size should match uploaded content size")
+}
+
+// TestMeasureFileSize_Integration_LargeFile tests file size measurement for larger files.
+func TestMeasureFileSize_Integration_LargeFile(t *testing.T) {
+	ctx := t.Context()
+
+	// Set up AWS container (localstack) and S3 client
+	cfg := &DocInitConfig{}
+	test.CreateAWSResources(t, cfg)
+
+	svc := New(cfg)
+
+	// Upload a larger test file (1MB)
+	bucket := cfg.GetBucket()
+	testKey := fmt.Sprintf("test/filesize/large/%s.bin", uuid.New().String())
+	expectedSize := int64(1024 * 1024) // 1MB
+	testContent := make([]byte, expectedSize)
+	for i := range testContent {
+		testContent[i] = byte(i % 256)
+	}
+
+	// Upload the test file to S3
+	_, err := cfg.GetStoreClient().PutObject(ctx, &s3.PutObjectInput{
+		Bucket: aws.String(bucket),
+		Key:    aws.String(testKey),
+		Body:   bytes.NewReader(testContent),
+	})
+	require.NoError(t, err, "failed to upload large test file to S3")
+
+	// Measure the file size using our function
+	measuredSize, err := svc.measureFileSize(ctx, bucket, testKey)
+	require.NoError(t, err, "measureFileSize should not return an error for large files")
+	require.NotNil(t, measuredSize, "measured size should not be nil")
+
+	// Verify the measured size matches expected
+	assert.Equal(t, expectedSize, *measuredSize, "measured file size should match 1MB")
+}
+
+// TestMeasureFileSize_Integration_NonExistentKey tests error handling for non-existent S3 keys.
+func TestMeasureFileSize_Integration_NonExistentKey(t *testing.T) {
+	ctx := t.Context()
+
+	// Set up AWS container (localstack) and S3 client
+	cfg := &DocInitConfig{}
+	test.CreateAWSResources(t, cfg)
+
+	svc := New(cfg)
+
+	bucket := cfg.GetBucket()
+	nonExistentKey := fmt.Sprintf("test/nonexistent/%s.txt", uuid.New().String())
+
+	// Try to measure size of non-existent file
+	_, err := svc.measureFileSize(ctx, bucket, nonExistentKey)
+	assert.Error(t, err, "measureFileSize should return an error for non-existent key")
+	assert.Contains(t, err.Error(), "HeadObject failed", "error should indicate HeadObject failure")
+}

internal/document/init/service.go

@@ -2,12 +2,14 @@ package documentinit
 
 import (
 	"queryorchestration/internal/serviceconfig"
+	"queryorchestration/internal/serviceconfig/objectstore"
 	documentsync "queryorchestration/internal/serviceconfig/queue/documentsync"
 )
 
 type ConfigProvider interface {
 	serviceconfig.ConfigProvider
 	documentsync.ConfigProvider
+	objectstore.ConfigProvider // Provides GetStoreClient() for S3 access to measure file size
 }
 
 type Service struct {

internal/document/init/service_test.go

@@ -4,6 +4,7 @@ import (
 	"testing"
 
 	"queryorchestration/internal/serviceconfig"
+	"queryorchestration/internal/serviceconfig/objectstore"
 	"queryorchestration/internal/serviceconfig/queue/documentsync"
 
 	"github.com/stretchr/testify/assert"
@@ -12,6 +13,7 @@ import (
 type DocInitConfig struct {
 	serviceconfig.BaseConfig
 	documentsync.DocSyncConfig
+	objectstore.ObjectStoreConfig
 }
 
 func TestNewDocumentInitService(t *testing.T) {

internal/document/service.go

@@ -36,6 +36,7 @@ type DocumentEnriched struct {
 	OriginalPath  *string
 	Labels        []LabelRecord
 	TextRecord    *TextRecord // Only populated when includeTextRecord=true
+	FileSizeBytes *int64      // File size in bytes (nil for legacy documents)
 }
 
 // LabelRecord represents a label applied to a document.

internal/folder/service.go

@@ -248,6 +248,7 @@ type DocumentInFolder struct {
 	OriginalPath  *string
 	Labels        []LabelRecord
 	TextRecord    *TextRecord // Only populated when includeTextRecord=true
+	FileSizeBytes *int64      // File size in bytes (nil for legacy documents)
 }
 
 // LabelRecord represents a label applied to a document.
@@ -368,6 +369,7 @@ func (s *Service) GetDocumentsEnriched(ctx context.Context, folderID uuid.UUID,
 			OriginalPath:  doc.Originalpath,
 			Labels:        labels,
 			TextRecord:    textRecordsMap[doc.ID],
+			FileSizeBytes: doc.FileSizeBytes,
 		}
 	}
 

internal/test/config.go

@@ -0,0 +1,152 @@
+// Package test provides configuration utilities for test infrastructure.
+// This file centralizes timeout configuration for CI/CD stability.
+package test
+
+import (
+	"fmt"
+	"log/slog"
+	"net/http"
+	"os"
+	"time"
+)
+
+// GetHTTPTimeout returns the HTTP client timeout for API calls.
+// Default is 30 seconds, configurable via E2E_HTTP_TIMEOUT environment variable.
+// In CI environments, this timeout ensures that stalled HTTP calls don't block
+// indefinitely, which was a root cause of test hangs.
+func GetHTTPTimeout() time.Duration {
+	if t := os.Getenv("E2E_HTTP_TIMEOUT"); t != "" {
+		if d, err := time.ParseDuration(t); err == nil {
+			return d
+		}
+	}
+	return 30 * time.Second
+}
+
+// GetPollingTimeout returns the timeout for polling operations like waiting
+// for client status or mock endpoint verification.
+// Default is 60 seconds locally, 180 seconds in CI (detected via CI or
+// BITBUCKET_PIPELINE_UUID environment variables).
+// Configurable via E2E_WAIT_TIMEOUT environment variable.
+func GetPollingTimeout() time.Duration {
+	if t := os.Getenv("E2E_WAIT_TIMEOUT"); t != "" {
+		if d, err := time.ParseDuration(t); err == nil {
+			return d
+		}
+	}
+	if isCI() {
+		return 180 * time.Second // 3 minutes in CI
+	}
+	return 60 * time.Second // 1 minute locally
+}
+
+// GetPollInterval returns the interval between polling attempts.
+// Default is 500ms, configurable via E2E_POLL_INTERVAL environment variable.
+func GetPollInterval() time.Duration {
+	if t := os.Getenv("E2E_POLL_INTERVAL"); t != "" {
+		if d, err := time.ParseDuration(t); err == nil {
+			return d
+		}
+	}
+	return 500 * time.Millisecond
+}
+
+// GetHTTPClient returns an HTTP client with the configured timeout.
+// This should be used instead of &http.Client{} to ensure all HTTP operations
+// have bounded timeouts and won't hang indefinitely.
+func GetHTTPClient() *http.Client {
+	return &http.Client{
+		Timeout: GetHTTPTimeout(),
+	}
+}
+
+// isCI returns true if running in a CI environment.
+// Detects CI via the CI or BITBUCKET_PIPELINE_UUID environment variables.
+func isCI() bool {
+	return os.Getenv("CI") != "" || os.Getenv("BITBUCKET_PIPELINE_UUID") != ""
+}
+
+// WaitForExternalReadiness waits for a service to be reachable from outside Docker.
+// This addresses the race condition where a container's internal health check passes
+// but the external port mapping isn't yet ready. This is particularly important in
+// Docker-in-Docker (DinD) environments like Bitbucket Pipelines.
+//
+// Parameters:
+//   - uri: The base URI of the service (e.g., "http://localhost:32768")
+//   - healthPath: The health check endpoint path (e.g., "/health" or "/liveness/probe")
+//   - timeout: Maximum time to wait for the service to become reachable
+//
+// Returns an error if the service is not reachable within the timeout period.
+func WaitForExternalReadiness(uri string, healthPath string, timeout time.Duration) error {
+	client := &http.Client{Timeout: 2 * time.Second}
+	deadline := time.Now().Add(timeout)
+
+	fullURL := fmt.Sprintf("%s%s", uri, healthPath)
+	slog.Info("Waiting for external readiness", "url", fullURL, "timeout", timeout)
+
+	var lastErr error
+	attempts := 0
+	for time.Now().Before(deadline) {
+		attempts++
+		resp, err := client.Get(fullURL)
+		if err == nil && resp.StatusCode == http.StatusOK {
+			resp.Body.Close()
+			slog.Info("Service is externally reachable", "url", fullURL, "attempts", attempts)
+			return nil
+		}
+		if err != nil {
+			lastErr = err
+		} else {
+			lastErr = fmt.Errorf("unexpected status code: %d", resp.StatusCode)
+			resp.Body.Close()
+		}
+		time.Sleep(500 * time.Millisecond)
+	}
+
+	return fmt.Errorf("service not reachable from test client after %v (%d attempts): %w", timeout, attempts, lastErr)
+}
+
+// WaitForAPIRoutes waits for API routes to be fully registered.
+// This addresses a race condition where the health endpoint returns 200 but other
+// API routes haven't been registered yet, causing "404 page not found" errors.
+// This is particularly important in slow CI environments where route registration
+// may take longer after the health endpoint becomes available.
+//
+// Parameters:
+//   - uri: The base URI of the service (e.g., "http://localhost:32768")
+//   - timeout: Maximum time to wait for routes to be registered
+//
+// Returns an error if routes are not available within the timeout period.
+func WaitForAPIRoutes(uri string, timeout time.Duration) error {
+	client := &http.Client{Timeout: 2 * time.Second}
+	deadline := time.Now().Add(timeout)
+
+	// Check /query endpoint - should return 200 with empty array when routes are ready
+	// A "404 page not found" means routes aren't registered yet
+	routeURL := fmt.Sprintf("%s/query", uri)
+	slog.Info("Waiting for API routes to be registered", "url", routeURL, "timeout", timeout)
+
+	var lastErr error
+	attempts := 0
+	for time.Now().Before(deadline) {
+		attempts++
+		resp, err := client.Get(routeURL)
+		if err == nil {
+			// Any response other than connection error means routes are registered
+			// 200 = success, 401 = auth required (but route exists), etc.
+			// Only "404 page not found" indicates routes not yet registered
+			if resp.StatusCode != http.StatusNotFound {
+				resp.Body.Close()
+				slog.Info("API routes are registered", "url", routeURL, "attempts", attempts, "status", resp.StatusCode)
+				return nil
+			}
+			lastErr = fmt.Errorf("routes not registered yet: %d", resp.StatusCode)
+			resp.Body.Close()
+		} else {
+			lastErr = err
+		}
+		time.Sleep(500 * time.Millisecond)
+	}
+
+	return fmt.Errorf("API routes not registered after %v (%d attempts): %w", timeout, attempts, lastErr)
+}

internal/test/config_test.go

@@ -0,0 +1,219 @@
+package test
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestGetHTTPTimeout(t *testing.T) {
+	// Test default value
+	timeout := GetHTTPTimeout()
+	assert.Equal(t, 30*time.Second, timeout, "default HTTP timeout should be 30 seconds")
+
+	// Test with environment variable
+	t.Setenv("E2E_HTTP_TIMEOUT", "45s")
+	timeout = GetHTTPTimeout()
+	assert.Equal(t, 45*time.Second, timeout, "HTTP timeout should respect E2E_HTTP_TIMEOUT env var")
+
+	// Test with invalid environment variable (should fall back to default)
+	t.Setenv("E2E_HTTP_TIMEOUT", "invalid")
+	timeout = GetHTTPTimeout()
+	assert.Equal(t, 30*time.Second, timeout, "invalid E2E_HTTP_TIMEOUT should fall back to default")
+}
+
+func TestGetPollingTimeout(t *testing.T) {
+	// Clear CI env vars to ensure we test non-CI default first
+	t.Setenv("CI", "")
+	t.Setenv("BITBUCKET_PIPELINE_UUID", "")
+
+	// Test default value (non-CI)
+	timeout := GetPollingTimeout()
+	assert.Equal(t, 60*time.Second, timeout, "default polling timeout should be 60 seconds")
+
+	// Test with environment variable override
+	t.Setenv("E2E_WAIT_TIMEOUT", "120s")
+	timeout = GetPollingTimeout()
+	assert.Equal(t, 120*time.Second, timeout, "polling timeout should respect E2E_WAIT_TIMEOUT env var")
+
+	// Test CI detection
+	t.Setenv("E2E_WAIT_TIMEOUT", "")
+	t.Setenv("CI", "true")
+	timeout = GetPollingTimeout()
+	assert.Equal(t, 180*time.Second, timeout, "polling timeout in CI should be 180 seconds")
+
+	// Test BITBUCKET_PIPELINE_UUID detection
+	t.Setenv("CI", "")
+	t.Setenv("BITBUCKET_PIPELINE_UUID", "some-uuid")
+	timeout = GetPollingTimeout()
+	assert.Equal(t, 180*time.Second, timeout, "polling timeout with BITBUCKET_PIPELINE_UUID should be 180 seconds")
+}
+
+func TestGetPollInterval(t *testing.T) {
+	// Test default value
+	interval := GetPollInterval()
+	assert.Equal(t, 500*time.Millisecond, interval, "default poll interval should be 500ms")
+
+	// Test with environment variable
+	t.Setenv("E2E_POLL_INTERVAL", "1s")
+	interval = GetPollInterval()
+	assert.Equal(t, 1*time.Second, interval, "poll interval should respect E2E_POLL_INTERVAL env var")
+
+	// Test with invalid environment variable (should fall back to default)
+	t.Setenv("E2E_POLL_INTERVAL", "invalid")
+	interval = GetPollInterval()
+	assert.Equal(t, 500*time.Millisecond, interval, "invalid E2E_POLL_INTERVAL should fall back to default")
+}
+
+func TestGetHTTPClient(t *testing.T) {
+	client := GetHTTPClient()
+	assert.NotNil(t, client, "GetHTTPClient should return a non-nil client")
+	assert.Equal(t, GetHTTPTimeout(), client.Timeout, "HTTP client timeout should match GetHTTPTimeout()")
+}
+
+func TestIsCI(t *testing.T) {
+	// Clear CI env vars to ensure consistent test behavior
+	t.Setenv("CI", "")
+	t.Setenv("BITBUCKET_PIPELINE_UUID", "")
+
+	// Test non-CI environment
+	assert.False(t, isCI(), "isCI should return false when CI env vars are not set")
+
+	// Test with CI env var
+	t.Setenv("CI", "true")
+	assert.True(t, isCI(), "isCI should return true when CI is set")
+
+	// Test with BITBUCKET_PIPELINE_UUID
+	t.Setenv("CI", "")
+	t.Setenv("BITBUCKET_PIPELINE_UUID", "some-uuid")
+	assert.True(t, isCI(), "isCI should return true when BITBUCKET_PIPELINE_UUID is set")
+}
+
+func TestWaitForExternalReadiness(t *testing.T) {
+	t.Run("success on first attempt", func(t *testing.T) {
+		// Create a test server that returns 200 OK
+		server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			if r.URL.Path == "/health" {
+				w.WriteHeader(http.StatusOK)
+				return
+			}
+			w.WriteHeader(http.StatusNotFound)
+		}))
+		defer server.Close()
+
+		err := WaitForExternalReadiness(server.URL, "/health", 5*time.Second)
+		assert.NoError(t, err, "should succeed when server is immediately ready")
+	})
+
+	t.Run("success after delay", func(t *testing.T) {
+		attempts := 0
+		server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			if r.URL.Path == "/health" {
+				attempts++
+				if attempts < 3 {
+					w.WriteHeader(http.StatusServiceUnavailable)
+					return
+				}
+				w.WriteHeader(http.StatusOK)
+				return
+			}
+			w.WriteHeader(http.StatusNotFound)
+		}))
+		defer server.Close()
+
+		err := WaitForExternalReadiness(server.URL, "/health", 5*time.Second)
+		assert.NoError(t, err, "should succeed after retries")
+		assert.GreaterOrEqual(t, attempts, 3, "should have made multiple attempts")
+	})
+
+	t.Run("timeout on non-responsive server", func(t *testing.T) {
+		// Create a server that always returns 503
+		server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			w.WriteHeader(http.StatusServiceUnavailable)
+		}))
+		defer server.Close()
+
+		err := WaitForExternalReadiness(server.URL, "/health", 2*time.Second)
+		assert.Error(t, err, "should fail when server never becomes ready")
+		assert.Contains(t, err.Error(), "not reachable", "error should indicate service not reachable")
+	})
+
+	t.Run("timeout on unreachable server", func(t *testing.T) {
+		// Use an invalid URL that will fail to connect
+		err := WaitForExternalReadiness("http://127.0.0.1:1", "/health", 2*time.Second)
+		assert.Error(t, err, "should fail when server is unreachable")
+		assert.Contains(t, err.Error(), "not reachable", "error should indicate service not reachable")
+	})
+}
+
+func TestWaitForAPIRoutes(t *testing.T) {
+	t.Run("success when routes return 200", func(t *testing.T) {
+		server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			if r.URL.Path == "/query" {
+				w.WriteHeader(http.StatusOK)
+				_, _ = w.Write([]byte("[]"))
+				return
+			}
+			w.WriteHeader(http.StatusNotFound)
+		}))
+		defer server.Close()
+
+		err := WaitForAPIRoutes(server.URL, 5*time.Second)
+		assert.NoError(t, err, "should succeed when routes return 200")
+	})
+
+	t.Run("success after routes become available", func(t *testing.T) {
+		attempts := 0
+		server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			if r.URL.Path == "/query" {
+				attempts++
+				if attempts < 3 {
+					// Simulate "404 page not found" before routes are registered
+					w.WriteHeader(http.StatusNotFound)
+					_, _ = w.Write([]byte("404 page not found"))
+					return
+				}
+				w.WriteHeader(http.StatusOK)
+				_, _ = w.Write([]byte("[]"))
+				return
+			}
+			w.WriteHeader(http.StatusNotFound)
+		}))
+		defer server.Close()
+
+		err := WaitForAPIRoutes(server.URL, 5*time.Second)
+		assert.NoError(t, err, "should succeed after routes become available")
+		assert.GreaterOrEqual(t, attempts, 3, "should have made multiple attempts")
+	})
+
+	t.Run("success when routes return 401", func(t *testing.T) {
+		// 401 means the route exists but requires auth - route is registered
+		server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			if r.URL.Path == "/query" {
+				w.WriteHeader(http.StatusUnauthorized)
+				return
+			}
+			w.WriteHeader(http.StatusNotFound)
+		}))
+		defer server.Close()
+
+		err := WaitForAPIRoutes(server.URL, 5*time.Second)
+		assert.NoError(t, err, "should succeed when routes return 401 (route exists)")
+	})
+
+	t.Run("timeout when routes never register", func(t *testing.T) {
+		// Server always returns 404 - routes never registered
+		server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			w.WriteHeader(http.StatusNotFound)
+			_, _ = w.Write([]byte("404 page not found"))
+		}))
+		defer server.Close()
+
+		err := WaitForAPIRoutes(server.URL, 2*time.Second)
+		assert.Error(t, err, "should fail when routes never register")
+		assert.Contains(t, err.Error(), "not registered", "error should indicate routes not registered")
+	})
+}

internal/test/ecosystem.go

@@ -70,7 +70,24 @@ func CreateFullNetwork(t testing.TB, ctx context.Context, cfg FullDependenciesCo
 
 	wg.Wait()
 
-	qService, err := queryapi.NewClientWithResponses(apiContainers[QueryAPIName].URI)
+	// Wait for external readiness before creating the client.
+	// This addresses the race condition where containers report healthy internally
+	// but the external port mapping isn't fully ready, causing 404 errors in CI.
+	err := WaitForExternalReadiness(apiContainers[QueryAPIName].URI, "/health", GetPollingTimeout())
+	require.NoError(t, err, "QueryAPI not externally reachable")
+
+	// Wait for API routes to be registered.
+	// The health endpoint may return 200 before all routes are registered,
+	// causing "404 page not found" errors on actual API calls.
+	err = WaitForAPIRoutes(apiContainers[QueryAPIName].URI, GetPollingTimeout())
+	require.NoError(t, err, "QueryAPI routes not registered")
+
+	// Create API client with bounded HTTP timeout to prevent indefinite hangs
+	// that were causing CI failures (goroutines blocking until 15-minute global timeout)
+	qService, err := queryapi.NewClientWithResponses(
+		apiContainers[QueryAPIName].URI,
+		queryapi.WithHTTPClient(GetHTTPClient()),
+	)
 	require.NoError(t, err)
 
 	return Network{

internal/test/mockserver.go

@@ -103,7 +103,7 @@ func CreateMockServer(t testing.TB) *MockServer {
 	require.NoError(t, err)
 
 	return &MockServer{
-		Client:    &http.Client{},
+		Client:    GetHTTPClient(),
 		Internal:  Address(fmt.Sprintf("http://%s:%d", mockServerAlias, port.Int())),
 		External:  Address(fmt.Sprintf("http://%s:%d", host, externalPort.Int())),
 		Container: container,
@@ -193,11 +193,12 @@ func WaitForMockEndpoint(t testing.TB, server *MockServer, request MockRequest)
 
 	verifyURL := fmt.Sprintf("%s/mockserver/verify", server.External)
 
-	clientTimeout := 500 * time.Millisecond
-	client := &http.Client{}
+	// Use configured HTTP client with timeout to prevent indefinite hangs
+	client := GetHTTPClient()
 
-	timeout := time.After(60 * time.Second)
-	ticker := time.NewTicker(clientTimeout)
+	// Use configurable polling timeout (longer in CI environments)
+	timeout := time.After(GetPollingTimeout())
+	ticker := time.NewTicker(GetPollInterval())
 	defer ticker.Stop()
 
 	slog.Info("Attempting to process request", "body", jsonData)
@@ -238,7 +239,8 @@ func retrieveMatchingRequest(t testing.TB, server *MockServer, request MockReque
 	req, err := http.NewRequest("PUT", retrieveURL, nil)
 	require.NoError(t, err)
 
-	client := &http.Client{}
+	// Use configured HTTP client with timeout to prevent indefinite hangs
+	client := GetHTTPClient()
 	resp, err := client.Do(req)
 	require.NoError(t, err)
 	defer resp.Body.Close()

internal/test/queryAPI/service.go

@@ -88,8 +88,10 @@ func SetQueryForClient(t testing.TB, client queryapi.ClientWithResponsesInterfac
 func WaitForClientStatus(t testing.TB, ctx context.Context, service queryapi.ClientWithResponsesInterface, id string, status queryapi.ClientStatus) {
 	t.Helper()
 
-	timeout := time.After(60 * time.Second)
-	ticker := time.NewTicker(500 * time.Millisecond) // Increased from 100ms to avoid rate limiting (2 req/s < 5 req/s limit)
+	// Use configurable polling timeout (longer in CI environments to handle slower infrastructure)
+	// and configurable poll interval to balance responsiveness with rate limiting concerns
+	timeout := time.After(test.GetPollingTimeout())
+	ticker := time.NewTicker(test.GetPollInterval())
 	defer ticker.Stop()
 
 	for {