chore: refactor.

container.go

@@ -203,14 +203,14 @@ func (t *Container) publish(route Route, opt ...PubOpt) (Intent, error) {
 		t.Size,
 		intents...,
 	)
+	if err != nil {
+		return nil, err
+	}
 	ir.OnClose(func() {
 		t.mu.Lock()
 		defer t.mu.Unlock()
 		delete(t.intentRouters, route.ID())
 	})
-	if err != nil {
-		return nil, err
-	}
 
 	t.intentRouters[route.ID()] = ir
 	return ir.Wrap(), nil
@@ -244,14 +244,14 @@ func (t *Container) subscribe(route Route, opt ...SubOpt) (Interest, error) {
 		t.Size,
 		interests...,
 	)
+	if err != nil {
+		return nil, err
+	}
 	ir.OnClose(func() {
 		t.mu.Lock()
 		defer t.mu.Unlock()
 		delete(t.interestRouters, route.ID())
 	})
-	if err != nil {
-		return nil, err
-	}
 
 	t.interestRouters[route.ID()] = ir
 	return ir.Wrap(), nil

container_test.go

@@ -121,6 +121,7 @@ func TestContainer_Publish(t *testing.T) {
 
 	mockIntent := new(MockIntent)
 	mockIntent.On("Route").Return(route)
+	mockIntent.On("OnClose", mock.Anything).Return(mockIntent)
 	ch := make(chan Route)
 	mockIntent.On("Interest").Return((<-chan Route)(ch))
 	endpoint := new(MockEndpoint)
@@ -168,6 +169,7 @@ func TestContainer_Subscribe(t *testing.T) {
 	mockInterest.On("Route").Return(route)
 	ch := make(chan proto.Message)
 	mockInterest.On("C").Return((<-chan proto.Message)(ch))
+	mockInterest.On("OnClose", mock.Anything).Return(mockInterest)
 	endpoint := new(MockEndpoint)
 	endpoint.On("Init", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil)
 	endpoint.On("Close").Return(nil)
@@ -213,11 +215,13 @@ func TestContainer_SubscribePublish(t *testing.T) {
 	mockIntent.On("Route").Return(route)
 	chI := make(chan Route, 1)
 	mockIntent.On("Interest").Return((<-chan Route)(chI))
+	mockIntent.On("OnClose", mock.Anything).Return(mockIntent)
 
 	mockInterest := new(MockInterestInternal)
 	mockInterest.On("Route").Return(route)
 	ch := make(chan proto.Message)
 	mockInterest.On("C").Return((<-chan proto.Message)(ch))
+	mockInterest.On("OnClose", mock.Anything).Return(mockInterest)
 	endpoint := new(MockEndpoint)
 	endpoint.On("Init", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil)
 	endpoint.On("Close").Return(nil)