Use protocol of latest owner

2 files changed, 8 insertions, 8 deletions

diff --git a/messagebus/src/main/java/com/yahoo/messagebus/network/NetworkMultiplexer.java b/messagebus/src/main/java/com/yahoo/messagebus/network/NetworkMultiplexer.java
index 4ebee093858..5e6a8c7688e 100644
--- a/messagebus/src/main/java/com/yahoo/messagebus/network/NetworkMultiplexer.java
+++ b/messagebus/src/main/java/com/yahoo/messagebus/network/NetworkMultiplexer.java
@@ -1,7 +1,6 @@
 // Copyright Verizon Media. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package com.yahoo.messagebus.network;
 
-import com.yahoo.component.AbstractComponent;
 import com.yahoo.messagebus.Message;
 import com.yahoo.messagebus.Protocol;
 import com.yahoo.text.Utf8Array;
@@ -81,12 +80,12 @@ public class NetworkMultiplexer implements NetworkOwner {
 
     @Override
     public Protocol getProtocol(Utf8Array name) {
-        for (NetworkOwner owner : owners) {
-            Protocol protocol = owner.getProtocol(name);
-            if (protocol != null)
-                return protocol;
-        }
-        return null;
+        // Should ideally couple this to the actual receiver ...
+        Protocol protocol = null;
+        for (NetworkOwner owner : owners)
+            protocol = owner.getProtocol(name) == null ? protocol : owner.getProtocol(name);
+
+        return protocol;
     }
 
     @Override
diff --git a/messagebus/src/test/java/com/yahoo/messagebus/network/NetworkMultiplexerTest.java b/messagebus/src/test/java/com/yahoo/messagebus/network/NetworkMultiplexerTest.java
index 544b4e84899..f1d83550601 100644
--- a/messagebus/src/test/java/com/yahoo/messagebus/network/NetworkMultiplexerTest.java
+++ b/messagebus/src/test/java/com/yahoo/messagebus/network/NetworkMultiplexerTest.java
@@ -60,8 +60,9 @@ public class NetworkMultiplexerTest {
         Protocol protocol1 = new SimpleProtocol();
         Protocol protocol2 = new SimpleProtocol();
         owner1.protocols.put(name, protocol1);
-        owner2.protocols.put(name, protocol2);
         assertEquals(protocol1, shared.getProtocol(name));
+        owner2.protocols.put(name, protocol2);
+        assertEquals(protocol2, shared.getProtocol(name));
 
         Message message1 = new SimpleMessage("one");
         Message message2 = new SimpleMessage("two");