package millproto
import (
"bytes"
"encoding/binary"
"strings"
"testing"
millv1 "tangled.org/core/spindle/mill/proto/gen"
"tangled.org/core/spindle/quota"
)
func TestEncodeDecodeRoundTrip(t *testing.T) {
var buf bytes.Buffer
enc := NewEncoder(&buf)
want := &Message{
ReserveSeat: &millv1.ReserveSeat{
LeaseId: "lease-1",
TargetEngine: "microvm",
RawWorkflowJson: `{"name":"build"}`,
PipelineId: "abc123",
TtlSeconds: 30,
RepoDid: "did:web:example.com",
},
}
if err := enc.Encode(want); err != nil {
t.Fatalf("Encode() error = %v", err)
}
got, err := NewDecoder(&buf).Decode()
if err != nil {
t.Fatalf("Decode() error = %v", err)
}
rs := got.GetReserveSeat()
if rs == nil {
t.Fatal("decoded message missing reserve_seat")
}
if rs.LeaseId != "lease-1" || rs.TargetEngine != "microvm" || rs.PipelineId != "abc123" || rs.TtlSeconds != 30 || rs.RepoDid != "did:web:example.com" {
t.Fatalf("round-trip mismatch: %+v", rs)
}
}
func TestDecoderRejectsOversizedMessage(t *testing.T) {
var tooLarge bytes.Buffer
var header [4]byte
binary.BigEndian.PutUint32(header[:], MaxMessageBytes+1)
tooLarge.Write(header[:])
if _, err := NewDecoder(&tooLarge).Decode(); err == nil {
t.Fatal("expected oversized message error")
}
}
func TestValidationRules(t *testing.T) {
tests := []struct {
name string
msg *Message
wantErr bool
}{
{
name: "valid ack message",
msg: &Message{
Ack: &millv1.Ack{
Epoch: "inc-1",
UpToSeqno: 5,
},
},
wantErr: false,
},
{
name: "valid hello message",
msg: &Message{
Hello: &millv1.Hello{ProtocolVersion: 1,
Arch: "amd64",
Labels: []string{"linux"},
Epoch: "inc-1",
CacheNamespace: "tar-zstd-v1/linux/amd64",
},
},
wantErr: false,
},
{
name: "invalid message with zero payloads",
msg: &Message{},
wantErr: true,
},
{
name: "invalid message with multiple payloads",
msg: &Message{
Ack: &millv1.Ack{Epoch: "inc-1", UpToSeqno: 5},
Committed: &millv1.Committed{LeaseId: "x"},
},
wantErr: true,
},
{
name: "invalid ack message missing epoch",
msg: &Message{
Ack: &millv1.Ack{
UpToSeqno: 5,
},
},
wantErr: true,
},
{
name: "invalid node snapshot with zero seqno",
msg: &Message{
NodeSnapshot: &millv1.NodeSnapshot{
Seqno: 0,
},
},
wantErr: true,
},
{
name: "valid node snapshot with positive seqno",
msg: &Message{
NodeSnapshot: &millv1.NodeSnapshot{
Seqno: 1,
},
},
wantErr: false,
},
{
name: "invalid stream batch with zero seqno entry",
msg: &Message{
EventBatch: &millv1.EventBatch{
Epoch: "inc-1",
Events: []*millv1.Event{
{
Seqno: 0,
LeaseId: "lease-1",
Payload: &millv1.Event_StatusEvent{
StatusEvent: &millv1.StatusEvent{
Status: millv1.NonterminalStatus_NONTERMINAL_STATUS_RUNNING,
},
},
},
},
},
},
wantErr: true,
},
{
name: "invalid stream batch with empty entries",
msg: &Message{
EventBatch: &millv1.EventBatch{
Epoch: "inc-1",
Events: []*millv1.Event{},
},
},
wantErr: true,
},
{
name: "invalid reserve result with unknown enum",
msg: &Message{
ReserveResult: &millv1.ReserveResult{
LeaseId: "lease-1",
RejectClass: millv1.RejectClass(99),
},
},
wantErr: true,
},
{
name: "invalid stream entry with malformed oneof (empty payload)",
msg: &Message{
EventBatch: &millv1.EventBatch{
Epoch: "inc-1",
Events: []*millv1.Event{
{
Seqno: 1,
LeaseId: "lease-1",
Payload: nil,
},
},
},
},
wantErr: true,
},
{
name: "valid stream batch status event",
msg: &Message{
EventBatch: &millv1.EventBatch{
Epoch: "inc-1",
Events: []*millv1.Event{
{
Seqno: 1,
LeaseId: "lease-1",
Payload: &millv1.Event_StatusEvent{
StatusEvent: &millv1.StatusEvent{
Status: millv1.NonterminalStatus_NONTERMINAL_STATUS_RUNNING,
},
},
},
},
},
},
wantErr: false,
},
{
name: "valid stream batch attempt result",
msg: &Message{
EventBatch: &millv1.EventBatch{
Epoch: "inc-1",
Events: []*millv1.Event{
{
Seqno: 1,
LeaseId: "lease-1",
Payload: &millv1.Event_AttemptResult{
AttemptResult: &millv1.AttemptResult{
Status: millv1.TerminalStatus_TERMINAL_STATUS_SUCCESS,
},
},
},
},
},
},
wantErr: false,
},
{
name: "valid quota reserve request",
msg: &Message{
QuotaReq: &millv1.QuotaRequest{
RequestId: "req-1",
Operation: millv1.QuotaOperation_QUOTA_OPERATION_RESERVE,
LeaseId: "lease-1",
Kind: "nix_cache",
Key: "key-1",
Resources: map[string]int64{"cache_storage_bytes": 100},
},
},
wantErr: false,
},
{
name: "invalid quota request - empty request_id",
msg: &Message{
QuotaReq: &millv1.QuotaRequest{
RequestId: "",
Operation: millv1.QuotaOperation_QUOTA_OPERATION_RESERVE,
LeaseId: "lease-1",
},
},
wantErr: true,
},
{
name: "invalid quota request - unspecified operation",
msg: &Message{
QuotaReq: &millv1.QuotaRequest{
RequestId: "req-1",
Operation: millv1.QuotaOperation_QUOTA_OPERATION_UNSPECIFIED,
LeaseId: "lease-1",
},
},
wantErr: true,
},
{
name: "invalid quota request - empty lease_id",
msg: &Message{
QuotaReq: &millv1.QuotaRequest{
RequestId: "req-1",
Operation: millv1.QuotaOperation_QUOTA_OPERATION_RESERVE,
LeaseId: "",
},
},
wantErr: true,
},
{
name: "invalid quota request - negative resource amount",
msg: &Message{
QuotaReq: &millv1.QuotaRequest{
RequestId: "req-1",
Operation: millv1.QuotaOperation_QUOTA_OPERATION_RESERVE,
LeaseId: "lease-1",
Resources: map[string]int64{"cache_storage_bytes": -1},
},
},
wantErr: true,
},
{
name: "valid quota transition request with reservation_id",
msg: &Message{
QuotaReq: &millv1.QuotaRequest{
RequestId: "req-1",
Operation: millv1.QuotaOperation_QUOTA_OPERATION_BEGIN_COMMIT,
LeaseId: "lease-1",
ReservationId: "res-1",
},
},
wantErr: false,
},
{
name: "valid reserve result reporting workflow resources",
msg: &Message{
ReserveResult: &millv1.ReserveResult{
LeaseId: "lease-1",
Accepted: true,
QuotaResources: map[string]int64{
"workflows": 1,
"memory_mib": 2048,
"vcpus": 2,
"disk_mib": 8192,
},
},
},
wantErr: false,
},
{
name: "valid reserve result reporting custom resources",
msg: &Message{
ReserveResult: &millv1.ReserveResult{
LeaseId: "lease-1",
Accepted: true,
QuotaResources: map[string]int64{
"gpu_count": 4,
},
},
},
wantErr: false,
},
{
name: "invalid reserve result - negative resource amount",
msg: &Message{
ReserveResult: &millv1.ReserveResult{
LeaseId: "lease-1",
Accepted: true,
QuotaResources: map[string]int64{"memory_mib": -1},
},
},
wantErr: true,
},
{
name: "valid quota response with resource field shape",
msg: &Message{
QuotaResp: &millv1.QuotaResponse{
RequestId: "req-1",
ReservationId: "res-1",
Allowed: true,
Resource: "cache_storage_bytes",
},
},
wantErr: false,
},
{
name: "invalid quota response - empty request_id",
msg: &Message{
QuotaResp: &millv1.QuotaResponse{
RequestId: "",
},
},
wantErr: true,
},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
err := validator.Validate(tc.msg)
if (err != nil) != tc.wantErr {
t.Fatalf("Validate() error = %v, wantErr = %v", err, tc.wantErr)
}
})
}
}
func TestQuotaResourceMapBounds(t *testing.T) {
tooMany := make(map[string]int64)
for i := range quota.MaxResourcePairs + 1 {
tooMany[strings.Repeat("x", i+1)] = 1
}
for _, resources := range []map[string]int64{
{"": 1},
{strings.Repeat("é", 33): 1},
{"large": quota.MaxResourceAmount + 1},
tooMany,
} {
msg := &Message{QuotaReq: &millv1.QuotaRequest{
RequestId: "req-1",
Operation: millv1.QuotaOperation_QUOTA_OPERATION_RESERVE,
LeaseId: "lease-1",
Kind: "custom",
Key: "key",
Resources: resources,
}}
if err := validator.Validate(msg); err == nil {
t.Fatalf("expected invalid resources: %v", resources)
}
}
}
func TestNixCacheMessageBounds(t *testing.T) {
hashes := func(n int) []string {
out := make([]string, n)
for i := range out {
out[i] = "hash"
}
return out
}
used := func(hashes []string) *Message {
return &Message{NixCacheUsed: &millv1.NixCacheUsed{LeaseId: "lease-1", Cache: "cache", StoreHashes: hashes}}
}
reserve := func(obj *millv1.NixObject) *Message {
return &Message{QuotaReq: &millv1.QuotaRequest{
RequestId: "req-1",
Operation: millv1.QuotaOperation_QUOTA_OPERATION_RESERVE,
LeaseId: "lease-1",
Kind: "nix_cache",
Key: "cache/abc",
Resources: map[string]int64{"cache_storage_bytes": 100},
NixObject: obj,
}}
}
object := func(refs []string) *millv1.NixObject {
return &millv1.NixObject{Cache: "cache", StoreHash: "abc", StorePath: "/nix/store/abc-x", NarSize: 100, References: refs}
}
for name, tc := range map[string]struct {
msg *Message
wantErr bool
}{
"use": {used(hashes(512)), false},
"use without hashes": {used(nil), true},
"use over the batch": {used(hashes(513)), true},
"object": {reserve(object(hashes(1024))), false},
"object too many refs": {reserve(object(hashes(1025))), true},
"object negative size": {reserve(&millv1.NixObject{Cache: "cache", StoreHash: "abc", NarSize: -1}), true},
} {
if err := validator.Validate(tc.msg); (err != nil) != tc.wantErr {
t.Errorf("%s: err = %v, wantErr %v", name, err, tc.wantErr)
}
}
}
func TestProtocolVersion(t *testing.T) {
if ProtocolVersion != 5 {
t.Fatalf("ProtocolVersion = %d, want 5", ProtocolVersion)
}
}