github.com/asynkron/protoactor-go@v0.0.0-20240308120642-ef91a6abee75/router/routeractor_pool_test.go (about)

     1  package router
     2  
     3  import (
     4  	"testing"
     5  
     6  	"github.com/asynkron/protoactor-go/actor"
     7  	"github.com/stretchr/testify/mock"
     8  )
     9  
    10  func TestPoolRouterActor_Receive_AddRoute(t *testing.T) {
    11  	state := new(testRouterState)
    12  
    13  	a := poolRouterActor{state: state}
    14  
    15  	p1 := system.NewLocalPID("p1")
    16  	c := new(mockContext)
    17  	c.On("Message").Return(&AddRoutee{PID: p1})
    18  	c.On("Watch", p1).Once()
    19  
    20  	state.On("GetRoutees").Return(&actor.PIDSet{})
    21  	state.On("SetRoutees", actor.NewPIDSet(p1)).Once()
    22  
    23  	a.Receive(c)
    24  	mock.AssertExpectationsForObjects(t, state, c)
    25  }
    26  
    27  func TestPoolRouterActor_Receive_AddRoute_NoDuplicates(t *testing.T) {
    28  	state := new(testRouterState)
    29  
    30  	a := poolRouterActor{state: state}
    31  
    32  	p1 := system.NewLocalPID("p1")
    33  	c := new(mockContext)
    34  	c.On("Message").Return(&AddRoutee{PID: p1})
    35  
    36  	state.On("GetRoutees").Return(actor.NewPIDSet(p1))
    37  
    38  	a.Receive(c)
    39  	mock.AssertExpectationsForObjects(t, state, c)
    40  }
    41  
    42  func TestPoolRouterActor_Receive_RemoveRoute(t *testing.T) {
    43  	state := new(testRouterState)
    44  
    45  	a := poolRouterActor{state: state}
    46  
    47  	p1, pr1 := spawnMockProcess("p1")
    48  	defer removeMockProcess(p1)
    49  	pr1.On("SendUserMessage", p1, &actor.PoisonPill{}).Once()
    50  
    51  	p2 := system.NewLocalPID("p2")
    52  	c := new(mockContext)
    53  	c.On("Message").Return(&RemoveRoutee{PID: p1})
    54  	c.On("Unwatch", p1).Once()
    55  
    56  	c.On("Send")
    57  
    58  	state.On("GetRoutees").Return(actor.NewPIDSet(p1, p2))
    59  	state.On("SetRoutees", actor.NewPIDSet(p2)).Once()
    60  
    61  	a.Receive(c)
    62  	mock.AssertExpectationsForObjects(t, state, c)
    63  }
    64  
    65  func TestPoolRouterActor_Receive_BroadcastMessage(t *testing.T) {
    66  	state := new(testRouterState)
    67  	a := poolRouterActor{state: state}
    68  
    69  	p1 := system.NewLocalPID("p1")
    70  	p2 := system.NewLocalPID("p2")
    71  
    72  	child := new(mockProcess)
    73  	child.On("SendUserMessage", mock.Anything, mock.Anything).Times(2)
    74  
    75  	system.ProcessRegistry.Add(child, "p1")
    76  	system.ProcessRegistry.Add(child, "p2")
    77  	defer func() {
    78  		system.ProcessRegistry.Remove(&actor.PID{Id: "p1"})
    79  		system.ProcessRegistry.Remove(&actor.PID{Id: "p2"})
    80  	}()
    81  
    82  	c := new(mockContext)
    83  	c.On("Message").Return(&BroadcastMessage{"hi"})
    84  	c.On("Sender").Return((*actor.PID)(nil))
    85  	c.On("RequestWithCustomSender").Twice()
    86  
    87  	state.On("GetRoutees").Return(actor.NewPIDSet(p1, p2))
    88  
    89  	a.Receive(c)
    90  	mock.AssertExpectationsForObjects(t, state, c, child)
    91  }