sdk/artifact/storage_test.go

package artifact

import (
	"bytes"
	"context"
	"crypto/sha256"
	"encoding/base64"
	"encoding/hex"
	"encoding/xml"
	"fmt"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync"
	"testing"
	"time"
)

func sha256Hex(data string) string {
	sum := sha256.Sum256([]byte(data))
	return hex.EncodeToString(sum[:])
}

func TestPushBlob(t *testing.T) {
	t.Parallel()

	const data = "helloworld" // 10 bytes, chunked into 4+4+2

	var mu sync.Mutex
	received := map[string][]byte{}
	var blockList []string

	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch r.URL.Query().Get("comp") {
		case "block":
			id := r.URL.Query().Get("blockid")
			// Delay the first block so it completes after later ones; the
			// committed block list must still be in ascending index order.
			raw, _ := base64.StdEncoding.DecodeString(id)
			if string(raw) == "000000" {
				time.Sleep(20 * time.Millisecond)
			}
			body, _ := io.ReadAll(r.Body)
			mu.Lock()
			received[id] = body
			mu.Unlock()
			w.WriteHeader(http.StatusCreated)
		case "blocklist":
			body, _ := io.ReadAll(r.Body)
			var bl struct {
				Latest []struct {
					Value string `xml:",chardata"`
				} `xml:"Latest"`
			}
			if err := xml.Unmarshal(body, &bl); err != nil {
				t.Errorf("invalid block list xml: %v", err)
			}
			for _, l := range bl.Latest {
				blockList = append(blockList, l.Value)
			}
			w.WriteHeader(http.StatusCreated)
		default:
			http.Error(w, "unexpected request", http.StatusBadRequest)
		}
	}))
	defer server.Close()

	res, err := PushBlob(context.Background(), nil, strings.NewReader(data), server.URL+"?sv=test", 4, 3)
	if err != nil {
		t.Fatal(err)
	}
	if res.Size != int64(len(data)) {
		t.Errorf("expected size %d, got %d", len(data), res.Size)
	}
	if res.SHA256Sum != sha256Hex(data) {
		t.Errorf("expected sha %q, got %q", sha256Hex(data), res.SHA256Sum)
	}

	if len(received) != 3 {
		t.Fatalf("expected 3 blocks, got %d", len(received))
	}

	wantChunks := []string{"hell", "owor", "ld"}
	for i, want := range wantChunks {
		id := base64.StdEncoding.EncodeToString([]byte(chunkID(i)))
		if got := string(received[id]); got != want {
			t.Errorf("block %d: expected %q, got %q", i, want, got)
		}
	}

	if len(blockList) != 3 {
		t.Fatalf("expected 3 block list entries, got %d", len(blockList))
	}
	for i, id := range blockList {
		raw, err := base64.StdEncoding.DecodeString(id)
		if err != nil {
			t.Fatalf("invalid block id %q: %v", id, err)
		}
		if got := string(raw); got != chunkID(i) {
			t.Errorf("block list position %d: expected %q, got %q", i, chunkID(i), got)
		}
	}
}

func TestPushBlob_Empty(t *testing.T) {
	t.Parallel()

	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusCreated)
	}))
	defer server.Close()

	if _, err := PushBlob(context.Background(), nil, strings.NewReader(""), server.URL+"?sv=test", 4, 3); err == nil {
		t.Fatal("expected error for empty input")
	}
}

func TestPushBlob_Error(t *testing.T) {
	t.Parallel()

	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		http.Error(w, "boom", http.StatusInternalServerError)
	}))
	defer server.Close()

	_, err := PushBlob(context.Background(), nil, strings.NewReader("helloworld"), server.URL+"?sv=test", 4, 3)
	if err == nil || !strings.Contains(err.Error(), "upload failed") {
		t.Fatalf("expected upload error, got %v", err)
	}
}

func TestPullBlob(t *testing.T) {
	t.Parallel()

	const data = "helloworld"
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		_, _ = io.WriteString(w, data)
	}))
	defer server.Close()

	var buf bytes.Buffer
	res, err := PullBlob(context.Background(), nil, &buf, server.URL)
	if err != nil {
		t.Fatal(err)
	}
	if got := buf.String(); got != data {
		t.Errorf("expected %q, got %q", data, got)
	}
	if res.Size != int64(len(data)) {
		t.Errorf("expected size %d, got %d", len(data), res.Size)
	}
	if res.SHA256Sum != sha256Hex(data) {
		t.Errorf("expected sha %q, got %q", sha256Hex(data), res.SHA256Sum)
	}
}

func TestPullBlob_Error(t *testing.T) {
	t.Parallel()

	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		http.Error(w, "nope", http.StatusNotFound)
	}))
	defer server.Close()

	if _, err := PullBlob(context.Background(), nil, io.Discard, server.URL); err == nil {
		t.Fatal("expected error for non-2xx response")
	}
}

func chunkID(i int) string {
	return fmt.Sprintf("%06d", i)
}