cre-1601: more tests; more coverage

Author: mchain0

pkg/workflows/ring/plugin_test.go

@@ -378,6 +378,113 @@ func TestPlugin_getHealthyShards(t *testing.T) {
 	}
 }
 
+func TestPlugin_NoHealthyShardsFallbackToShardZero(t *testing.T) {
+	lggr := logger.Test(t)
+	store := NewStore()
+
+	// Set all shards unhealthy - store starts in transition state
+	store.SetAllShardHealth(map[uint32]bool{0: false, 1: false, 2: false})
+
+	config := ocr3types.ReportingPluginConfig{
+		N: 4, F: 1,
+	}
+
+	plugin, err := NewPlugin(store, config, lggr, &ConsensusConfig{
+		BatchSize:  100,
+		TimeToSync: 1 * time.Second,
+	})
+	require.NoError(t, err)
+
+	transmitter := NewTransmitter(lggr, store, nil, "test-account")
+
+	ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond)
+	defer cancel()
+
+	// Start a goroutine that requests allocation (will block waiting for OCR)
+	resultCh := make(chan uint32)
+	errCh := make(chan error, 1)
+	go func() {
+		shard, err := store.GetShardForWorkflow(ctx, "workflow-123")
+		if err != nil {
+			errCh <- err
+			return
+		}
+		resultCh <- shard
+	}()
+
+	// Give goroutine time to enqueue request
+	time.Sleep(10 * time.Millisecond)
+
+	// Verify request is pending for OCR consensus
+	pending := store.GetPendingAllocations()
+	require.Contains(t, pending, "workflow-123")
+
+	// Simulate OCR round with observations showing no healthy shards
+	// The pending allocation "workflow-123" should be included in observation
+	now := time.Now()
+	aos := make([]types.AttributedObservation, 3)
+	for i := 0; i < 3; i++ {
+		pbObs := &pb.Observation{
+			ShardHealthStatus: map[uint32]bool{0: false, 1: false, 2: false},
+			WorkflowIds:       []string{"workflow-123"},
+			Now:               timestamppb.New(now),
+		}
+		rawObs, err := proto.Marshal(pbObs)
+		require.NoError(t, err)
+		aos[i] = types.AttributedObservation{
+			Observation: rawObs,
+			Observer:    commontypes.OracleID(i),
+		}
+	}
+
+	// Use a previous outcome in steady state so we can test the fallback
+	priorOutcome := &pb.Outcome{
+		State: &pb.RoutingState{
+			Id:    1,
+			State: &pb.RoutingState_RoutableShards{RoutableShards: 3},
+		},
+		Routes: map[string]*pb.WorkflowRoute{},
+	}
+	priorBytes, err := proto.Marshal(priorOutcome)
+	require.NoError(t, err)
+
+	outcomeCtx := ocr3types.OutcomeContext{
+		SeqNr:           2,
+		PreviousOutcome: priorBytes,
+	}
+
+	// Run plugin Outcome phase
+	outcome, err := plugin.Outcome(ctx, outcomeCtx, nil, aos)
+	require.NoError(t, err)
+
+	// Transmit the outcome (applies routes to store)
+	reports, err := plugin.Reports(ctx, 2, outcome)
+	require.NoError(t, err)
+	require.Len(t, reports, 1)
+
+	err = transmitter.Transmit(ctx, types.ConfigDigest{}, 2, reports[0].ReportWithInfo, nil)
+	require.NoError(t, err)
+
+	// Blocked goroutine should now receive result from OCR - should be shard 0 (fallback)
+	select {
+	case shard := <-resultCh:
+		require.Equal(t, uint32(0), shard, "should fallback to shard 0 when no healthy shards")
+	case err := <-errCh:
+		t.Fatalf("unexpected error: %v", err)
+	case <-time.After(100 * time.Millisecond):
+		t.Fatal("allocation was not fulfilled by OCR")
+	}
+
+	// Verify the outcome assigned workflow-123 to shard 0
+	outcomeProto := &pb.Outcome{}
+	err = proto.Unmarshal(outcome, outcomeProto)
+	require.NoError(t, err)
+
+	route, exists := outcomeProto.Routes["workflow-123"]
+	require.True(t, exists, "workflow-123 should be in routes")
+	require.Equal(t, uint32(0), route.Shard, "workflow-123 should be assigned to shard 0 (fallback)")
+}
+
 func TestPlugin_ObservationQuorum(t *testing.T) {
 	lggr := logger.Test(t)
 	store := NewStore()

pkg/workflows/ring/store.go

@@ -28,13 +28,15 @@ type Store struct {
 	mu sync.Mutex
 }
 
+const AllocationRequestChannelCapacity = 1000
+
 func NewStore() *Store {
 	return &Store{
 		routingState:  make(map[string]uint32),
 		shardHealth:   make(map[uint32]bool),
 		healthyShards: make([]uint32, 0),
 		pendingAllocs: make(map[string][]chan uint32),
-		allocRequests: make(chan AllocationRequest, 1000),
+		allocRequests: make(chan AllocationRequest, AllocationRequestChannelCapacity),
 		mu:            sync.Mutex{},
 	}
 }

pkg/workflows/ring/store_test.go

@@ -128,23 +128,6 @@ func TestStore_GetHealthyShards(t *testing.T) {
 	require.Equal(t, []uint32{1, 2, 3}, healthyShards)
 }
 
-func TestStore_NilHashRingFallback(t *testing.T) {
-	store := NewStore()
-	ctx := context.Background()
-
-	// Set all shards unhealthy
-	store.SetAllShardHealth(map[uint32]bool{0: false, 1: false, 2: false})
-	// Simulate OCR moved to steady (even with no healthy shards)
-	store.SetRoutingState(&pb.RoutingState{
-		State: &pb.RoutingState_RoutableShards{RoutableShards: 0},
-	})
-
-	// Should not panic, should return 0 as fallback (no healthy shards)
-	shard, err := store.GetShardForWorkflow(ctx, "workflow-123")
-	require.NoError(t, err)
-	require.Equal(t, uint32(0), shard)
-}
-
 func TestStore_DistributionAcrossShards(t *testing.T) {
 	store := NewStore()
 	ctx := context.Background()
@@ -186,6 +169,50 @@ func sum(distribution map[uint32]int) int {
 	return total
 }
 
+func TestStore_GetShardForWorkflow_CacheHit(t *testing.T) {
+	store := NewStore()
+	ctx := context.Background()
+
+	// Set up steady state
+	store.SetAllShardHealth(map[uint32]bool{0: true, 1: true, 2: true})
+	store.SetRoutingState(&pb.RoutingState{
+		State: &pb.RoutingState_RoutableShards{RoutableShards: 3},
+	})
+
+	// Pre-populate cache with a specific shard assignment
+	store.SetShardForWorkflow("cached-workflow", 2)
+
+	// Should return cached value, not recompute
+	shard, err := store.GetShardForWorkflow(ctx, "cached-workflow")
+	require.NoError(t, err)
+	require.Equal(t, uint32(2), shard)
+}
+
+func TestStore_GetShardForWorkflow_ContextCancelledDuringSend(t *testing.T) {
+	store := NewStore()
+
+	// Put store in transition state
+	store.SetAllShardHealth(map[uint32]bool{0: true})
+	store.SetRoutingState(&pb.RoutingState{
+		State: &pb.RoutingState_Transition{
+			Transition: &pb.Transition{WantShards: 2},
+		},
+	})
+
+	// Fill up the allocRequests channel
+	for i := 0; i < AllocationRequestChannelCapacity; i++ {
+		store.allocRequests <- AllocationRequest{WorkflowID: "filler"}
+	}
+
+	// Context that's already cancelled
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+
+	// Should fail: channel is full and context is cancelled
+	_, err := store.GetShardForWorkflow(ctx, "workflow-123")
+	require.ErrorIs(t, err, context.Canceled)
+}
+
 func TestStore_PendingAllocsDuringTransition(t *testing.T) {
 	store := NewStore()
 	store.SetAllShardHealth(map[uint32]bool{0: true, 1: true})

pkg/workflows/ring/transmitter_test.go

@@ -18,6 +18,7 @@ import (
 type mockArbiterScaler struct {
 	called  bool
 	nShards uint32
+	err     error
 }
 
 func (m *mockArbiterScaler) Status(ctx context.Context, in *emptypb.Empty, opts ...grpc.CallOption) (*pb.ReplicaStatus, error) {
@@ -27,6 +28,9 @@ func (m *mockArbiterScaler) Status(ctx context.Context, in *emptypb.Empty, opts
 func (m *mockArbiterScaler) ConsensusWantShards(ctx context.Context, req *pb.ConsensusWantShardsRequest, opts ...grpc.CallOption) (*emptypb.Empty, error) {
 	m.called = true
 	m.nShards = req.NShards
+	if m.err != nil {
+		return nil, m.err
+	}
 	return &emptypb.Empty{}, nil
 }
 
@@ -119,3 +123,51 @@ func TestTransmitter_Transmit_TransitionState(t *testing.T) {
 	require.NoError(t, err)
 	require.Equal(t, uint32(5), mock.nShards)
 }
+
+func TestTransmitter_Transmit_InvalidReport(t *testing.T) {
+	lggr := logger.Test(t)
+	store := NewStore()
+	tx := NewTransmitter(lggr, store, nil, "test-account")
+
+	// Send invalid protobuf data
+	report := ocr3types.ReportWithInfo[[]byte]{Report: []byte("invalid protobuf")}
+	err := tx.Transmit(context.Background(), types.ConfigDigest{}, 0, report, nil)
+	require.Error(t, err)
+}
+
+func TestTransmitter_Transmit_ArbiterError(t *testing.T) {
+	lggr := logger.Test(t)
+	store := NewStore()
+	mock := &mockArbiterScaler{err: context.DeadlineExceeded}
+	tx := NewTransmitter(lggr, store, mock, "test-account")
+
+	outcome := &pb.Outcome{
+		State: &pb.RoutingState{
+			Id:    1,
+			State: &pb.RoutingState_RoutableShards{RoutableShards: 3},
+		},
+	}
+	outcomeBytes, _ := proto.Marshal(outcome)
+
+	err := tx.Transmit(context.Background(), types.ConfigDigest{}, 0, ocr3types.ReportWithInfo[[]byte]{Report: outcomeBytes}, nil)
+	require.ErrorIs(t, err, context.DeadlineExceeded)
+}
+
+func TestTransmitter_Transmit_NilState(t *testing.T) {
+	lggr := logger.Test(t)
+	store := NewStore()
+	tx := NewTransmitter(lggr, store, nil, "test-account")
+
+	outcome := &pb.Outcome{
+		State:  nil,
+		Routes: map[string]*pb.WorkflowRoute{"wf-1": {Shard: 0}},
+	}
+	outcomeBytes, _ := proto.Marshal(outcome)
+
+	err := tx.Transmit(context.Background(), types.ConfigDigest{}, 0, ocr3types.ReportWithInfo[[]byte]{Report: outcomeBytes}, nil)
+	require.NoError(t, err)
+
+	// Routes should still be applied
+	routes := store.GetAllRoutingState()
+	require.Equal(t, uint32(0), routes["wf-1"])
+}