@@ -114,61 +114,90 @@ func TestPoolSize(t *testing.T) {
114114}
115115
116116func TestPoolPublish (t * testing.T ) {
117- poolSize := 3
118- mockClients := make ([]* mockMQTTClient , poolSize )
119- mockTokens := make ([]* mockToken , poolSize )
117+ tests := []struct {
118+ name string
119+ poolSize int
120+ publishCount int
121+ expectedCalls int
122+ }{
123+ {
124+ name : "Pool size 2 Publish 4" ,
125+ poolSize : 2 ,
126+ publishCount : 4 ,
127+ expectedCalls : 2 ,
128+ },
129+ {
130+ name : "Pool size 3 Publish 12" ,
131+ poolSize : 3 ,
132+ publishCount : 12 ,
133+ expectedCalls : 4 ,
134+ },
135+ {
136+ name : "Pool size 4 Publish 24" ,
137+ poolSize : 4 ,
138+ publishCount : 24 ,
139+ expectedCalls : 6 ,
140+ },
141+ }
120142
121- for i := 0 ; i < poolSize ; i ++ {
122- mockClients [i ] = newMockMQTTClient (t )
123- mockTokens [i ] = newMockToken (t )
143+ for _ , tt := range tests {
144+ t .Run (tt .name , func (t * testing.T ) {
145+ mockTokens := make ([]* mockToken , tt .poolSize )
146+ mockClients := make ([]* mockMQTTClient , tt .poolSize )
124147
125- mockTokens [ i ]. On ( "Wait" ). Return ( true )
126- mockTokens [i ]. On ( "WaitTimeout" , mock . Anything ). Return ( true )
127- mockTokens [i ]. On ( "Error" ). Return ( nil )
148+ for i := 0 ; i < tt . poolSize ; i ++ {
149+ mockClients [i ] = newMockMQTTClient ( t )
150+ mockTokens [i ] = newMockToken ( t )
128151
129- mockClients [i ].On ("IsConnectionOpen" ).Return (true )
130- mockClients [i ].On ("Connect" ).Return (mockTokens [i ])
131- mockClients [i ].On ("Disconnect" , mock .Anything )
132- mockClients [i ].On ("Publish" , mock .Anything , mock .Anything , mock .Anything , mock .Anything ).Return (mockTokens [i ])
133- }
152+ mockTokens [i ].On ("Wait" ).Return (true )
153+ mockTokens [i ].On ("WaitTimeout" , mock .Anything ).Return (true )
154+ mockTokens [i ].On ("Error" ).Return (nil )
134155
135- clientIndex := 0
136- originalFunc := newClientFunc .Load ()
137- newClientFunc .Store (func (opts * mqtt.ClientOptions ) mqtt.Client {
138- client := mockClients [clientIndex % poolSize ]
139- clientIndex ++
140- return client
141- })
142- defer newClientFunc .Store (originalFunc )
156+ mockClients [i ].On ("IsConnectionOpen" ).Return (true )
157+ mockClients [i ].On ("Connect" ).Return (mockTokens [i ])
158+ mockClients [i ].On ("Disconnect" , mock .Anything )
159+ mockClients [i ].On ("Publish" , mock .Anything , mock .Anything , mock .Anything , mock .Anything ).Return (mockTokens [i ])
160+ }
143161
144- client , err := NewClient (
145- WithAddress ("localhost" , 1883 ),
146- WithClientID ("test" ),
147- WithPoolSize (poolSize ),
148- )
149- assert .NoError (t , err )
150- assert .NotNil (t , client )
162+ clientIndex := 0
163+ originalFunc := newClientFunc .Load ()
164+ newClientFunc .Store (func (opts * mqtt.ClientOptions ) mqtt.Client {
165+ client := mockClients [clientIndex % tt .poolSize ]
166+ clientIndex ++
167+ return client
168+ })
169+ defer newClientFunc .Store (originalFunc )
151170
152- err = client .Start ()
153- assert .NoError (t , err )
171+ client , err := NewClient (
172+ WithAddress ("localhost" , 1883 ),
173+ WithClientID ("test" ),
174+ WithPoolSize (tt .poolSize ),
175+ )
176+ assert .NoError (t , err )
177+ assert .NotNil (t , client )
154178
155- for i := range poolSize {
156- atomic .StoreInt64 (& mockClients [i ].publishCallCount , 0 )
157- }
179+ err = client .Start ()
180+ assert .NoError (t , err )
158181
159- expectedPublishes := 6
160- for i := 0 ; i < expectedPublishes ; i ++ {
161- err = client .Publish (context .Background (), "topic" , []byte ("test" ))
162- assert .NoError (t , err )
163- }
182+ for i := range tt .poolSize {
183+ atomic .StoreInt64 (& mockClients [i ].publishCallCount , 0 )
184+ }
164185
165- totalPublishCount := int64 ( 0 )
166- for i := range poolSize {
167- totalPublishCount += atomic . LoadInt64 ( & mockClients [ i ]. publishCallCount )
168- }
186+ for i := 0 ; i < tt . publishCount ; i ++ {
187+ err = client . Publish ( context . Background (), "topic" , [] byte ( "test" ))
188+ assert . NoError ( t , err )
189+ }
169190
170- assert .Equal (t , int64 (expectedPublishes ), totalPublishCount , "Total publish count should match expected" )
171- assert .NoError (t , client .stop ())
191+ totalPublishCount := int64 (0 )
192+ for i := range tt .poolSize {
193+ totalPublishCount += atomic .LoadInt64 (& mockClients [i ].publishCallCount )
194+ assert .Equal (t , int64 (tt .expectedCalls ), atomic .LoadInt64 (& mockClients [i ].publishCallCount ), "Each client should have correct publish call count" )
195+ }
196+
197+ assert .Equal (t , int64 (tt .publishCount ), totalPublishCount , "Total publish count should match expected" )
198+ assert .NoError (t , client .stop ())
199+ })
200+ }
172201}
173202
174203func TestPoolSubscribe (t * testing.T ) {
0 commit comments