fixes WeightedLoadbalanceStrategy test

Signed-off-by: Oleh Dokuka <[email protected]>

Author: OlegDokuka

rsocket-core/src/test/java/io/rsocket/loadbalance/LoadbalanceTest.java

@@ -132,10 +132,14 @@ public void shouldDeliverAllTheRequestsWithWeightedStrategy() throws Interrupted
               source,
               WeightedLoadbalanceStrategy.builder()
                   .weightedStatsResolver(
-                      rsocket ->
-                          ((PooledRSocket) rsocket).target() == target1
-                              ? weightedRSocket1
-                              : weightedRSocket2)
+                      rsocket -> {
+                        if (rsocket instanceof TestRSocket) {
+                          return (WeightedRSocket) ((TestRSocket) rsocket).source();
+                        }
+                        return ((PooledRSocket) rsocket).target() == target1
+                            ? weightedRSocket1
+                            : weightedRSocket2;
+                      })
                   .build());
 
       RaceTestUtils.race(
@@ -315,6 +319,10 @@ public Mono<Void> onClose() {
     public void dispose() {
       sink.tryEmitEmpty();
     }
+
+    public RSocket source() {
+      return source;
+    }
   }
 
   private static class WeightedRSocket extends BaseWeightedStats implements RSocket {