[p2p/StreamManager] get stream tests

dominant-strategies · Sep 3, 2024 · 6507f76 · 6507f76
1 parent 2a447d7
commit 6507f76
Showing 1 changed file with 11 additions and 1 deletion.
diff --git a/p2p/node/streamManager/streamManager_test.go b/p2p/node/streamManager/streamManager_test.go
@@ -43,15 +43,25 @@ func TestOpenStream(t *testing.T) {
 	MockLibp2pStream := mock_p2p.NewMockStream(ctrl)
 	MockConn := mock_p2p.NewMockConn(ctrl)
 
-	mockNode.EXPECT().GetBandwidthCounter().Return(nil).Times(1)
+	mockNode.EXPECT().GetBandwidthCounter().Return(nil).AnyTimes()
 	MockLibp2pStream.EXPECT().Close().Return(nil).AnyTimes()
 	MockLibp2pStream.EXPECT().Conn().Return(MockConn).Times(1)
 	MockLibp2pStream.EXPECT().Protocol().Return(protocol.ProtocolVersion).Times(1)
 	MockLibp2pStream.EXPECT().Read(gomock.Any()).Return(0, nil).AnyTimes()
 	MockConn.EXPECT().RemotePeer().Return(peerID)
 	mockHost.EXPECT().NewStream(gomock.Any(), gomock.Any(), gomock.Any()).Return(MockLibp2pStream, nil).Times(1)
 
+	// GetStream Error
+	entry, err := sm.GetStream(peerID)
+	require.Error(t, err)
+	require.Nil(t, entry)
+
 	err = sm.OpenStream(mockHost2.ID())
 
 	require.NoError(t, err)
+
+	// Get Stream Success
+	entry, err = sm.GetStream(peerID)
+	require.NoError(t, err)
+	require.Equal(t, MockLibp2pStream, entry)
 }