at-container-registry/pkg/hold/oci/multipart_test.go

package oci

import (
	"testing"
	"time"
)

// Tests for MultipartManager

func TestCreateSession(t *testing.T) {
	mgr := &MultipartManager{
		sessions: make(map[string]*MultipartSession),
	}

	session := mgr.CreateSession("sha256:test123", Buffered, "")

	if session.UploadID == "" {
		t.Error("Expected non-empty uploadID")
	}
	if session.Digest != "sha256:test123" {
		t.Errorf("Expected digest=sha256:test123, got %s", session.Digest)
	}
	if session.Mode != Buffered {
		t.Errorf("Expected mode=Buffered, got %v", session.Mode)
	}
	if session.Parts == nil {
		t.Error("Expected Parts map to be initialized")
	}
	if session.CreatedAt.IsZero() {
		t.Error("Expected CreatedAt to be set")
	}
}

func TestCreateSession_S3Native(t *testing.T) {
	mgr := &MultipartManager{
		sessions: make(map[string]*MultipartSession),
	}

	s3UploadID := "aws-multipart-id-123"
	session := mgr.CreateSession("sha256:test123", S3Native, s3UploadID)

	if session.Mode != S3Native {
		t.Errorf("Expected mode=S3Native, got %v", session.Mode)
	}
	if session.S3UploadID != s3UploadID {
		t.Errorf("Expected S3UploadID=%s, got %s", s3UploadID, session.S3UploadID)
	}
}

func TestGetSession_Success(t *testing.T) {
	mgr := &MultipartManager{
		sessions: make(map[string]*MultipartSession),
	}

	created := mgr.CreateSession("sha256:test123", Buffered, "")

	retrieved, err := mgr.GetSession(created.UploadID)
	if err != nil {
		t.Fatalf("Expected success, got error: %v", err)
	}

	if retrieved.UploadID != created.UploadID {
		t.Errorf("Expected uploadID=%s, got %s", created.UploadID, retrieved.UploadID)
	}
}

func TestGetSession_NotFound(t *testing.T) {
	mgr := &MultipartManager{
		sessions: make(map[string]*MultipartSession),
	}

	_, err := mgr.GetSession("non-existent-id")
	if err == nil {
		t.Error("Expected error for non-existent session")
	}
}

func TestDeleteSession(t *testing.T) {
	mgr := &MultipartManager{
		sessions: make(map[string]*MultipartSession),
	}

	session := mgr.CreateSession("sha256:test123", Buffered, "")
	uploadID := session.UploadID

	// Verify it exists
	_, err := mgr.GetSession(uploadID)
	if err != nil {
		t.Fatalf("Session should exist before deletion")
	}

	// Delete it
	mgr.DeleteSession(uploadID)

	// Verify it's gone
	_, err = mgr.GetSession(uploadID)
	if err == nil {
		t.Error("Session should not exist after deletion")
	}
}

func TestStorePart(t *testing.T) {
	session := &MultipartSession{
		UploadID: "test-upload",
		Digest:   "sha256:test",
		Mode:     Buffered,
		Parts:    make(map[int]*MultipartPart),
	}

	data := []byte("test part data")
	etag := session.StorePart(1, data)

	if etag == "" {
		t.Error("Expected non-empty etag")
	}

	part, exists := session.Parts[1]
	if !exists {
		t.Fatal("Part 1 should exist")
	}

	if part.PartNumber != 1 {
		t.Errorf("Expected partNumber=1, got %d", part.PartNumber)
	}
	if string(part.Data) != string(data) {
		t.Errorf("Expected data=%s, got %s", string(data), string(part.Data))
	}
	if part.ETag != etag {
		t.Errorf("Expected etag=%s, got %s", etag, part.ETag)
	}
	if part.Size != int64(len(data)) {
		t.Errorf("Expected size=%d, got %d", len(data), part.Size)
	}
}

func TestAssembleBufferedParts_Success(t *testing.T) {
	session := &MultipartSession{
		UploadID: "test-upload",
		Digest:   "sha256:test",
		Mode:     Buffered,
		Parts:    make(map[int]*MultipartPart),
	}

	// Add parts in non-sequential order to test sorting
	session.StorePart(2, []byte("second part"))
	session.StorePart(1, []byte("first part"))
	session.StorePart(3, []byte("third part"))

	data, size, err := session.AssembleBufferedParts()
	if err != nil {
		t.Fatalf("Expected success, got error: %v", err)
	}

	expected := "first partsecond partthird part"
	if string(data) != expected {
		t.Errorf("Expected data=%s, got %s", expected, string(data))
	}

	if size != int64(len(expected)) {
		t.Errorf("Expected size=%d, got %d", len(expected), size)
	}
}

func TestAssembleBufferedParts_MissingPart(t *testing.T) {
	session := &MultipartSession{
		UploadID: "test-upload",
		Digest:   "sha256:test",
		Mode:     Buffered,
		Parts:    make(map[int]*MultipartPart),
	}

	// Add parts 1 and 3, but not 2
	session.StorePart(1, []byte("first part"))
	session.StorePart(3, []byte("third part"))

	_, _, err := session.AssembleBufferedParts()
	if err == nil {
		t.Error("Expected error for missing part 2")
	}
}

func TestAssembleBufferedParts_WrongMode(t *testing.T) {
	session := &MultipartSession{
		UploadID: "test-upload",
		Digest:   "sha256:test",
		Mode:     S3Native,
		Parts:    make(map[int]*MultipartPart),
	}

	_, _, err := session.AssembleBufferedParts()
	if err == nil {
		t.Error("Expected error for S3Native mode")
	}
}

func TestCleanupExpiredSessions(t *testing.T) {
	mgr := &MultipartManager{
		sessions: make(map[string]*MultipartSession),
	}

	// Create an old session (>24h)
	oldSession := &MultipartSession{
		UploadID:     "old-session",
		Digest:       "sha256:old",
		Mode:         Buffered,
		Parts:        make(map[int]*MultipartPart),
		CreatedAt:    time.Now().Add(-25 * time.Hour),
		LastActivity: time.Now().Add(-25 * time.Hour),
	}
	mgr.sessions[oldSession.UploadID] = oldSession

	// Create a recent session
	recentSession := mgr.CreateSession("sha256:recent", Buffered, "")

	// Run cleanup
	mgr.cleanupExpiredSessions()

	// Old session should be gone
	_, err := mgr.GetSession("old-session")
	if err == nil {
		t.Error("Old session should have been cleaned up")
	}

	// Recent session should still exist
	_, err = mgr.GetSession(recentSession.UploadID)
	if err != nil {
		t.Error("Recent session should still exist")
	}
}