Merge pull request #576 from sasvaritoni/master

Fix PredictiveUnitState image name and version (#562)
SeldonIO · May 18, 2019 · e03b4e8 · e03b4e8
2 parents 83bdb2d + 94cbc61
commit e03b4e8
Show file tree

Hide file tree

Showing 2 changed files with 43 additions and 4 deletions.
diff --git a/engine/src/main/java/io/seldon/engine/predictors/PredictiveUnitState.java b/engine/src/main/java/io/seldon/engine/predictors/PredictiveUnitState.java
@@ -87,11 +87,11 @@ public PredictiveUnitState(
 
 		if (containersMap.containsKey(name)){
 			this.image = containersMap.get(name).getImage();
-			if (image.contains(":"))
+			int i = image.lastIndexOf(":");
+			if (i >= 0)
 			{
-				String[] parts = image.split(":");
-				this.imageName = parts[0];
-				this.imageVersion = parts[1];
+				this.imageName = StringUtils.substring(image, 0, i);
+				this.imageVersion = StringUtils.substring(image, i+1);
 			}
 			else
 			{

diff --git a/engine/src/test/java/io/seldon/engine/predictors/PredictiveUnitStateTest.java b/engine/src/test/java/io/seldon/engine/predictors/PredictiveUnitStateTest.java
@@ -0,0 +1,39 @@
+package io.seldon.engine.predictors;
+
+import java.util.HashMap;
+import java.util.Map;
+
+import org.junit.Assert;
+import org.junit.Test;
+
+import io.kubernetes.client.proto.V1.Container;
+import io.seldon.protos.DeploymentProtos.Endpoint;
+import io.seldon.protos.DeploymentProtos.PredictiveUnit;
+
+public class PredictiveUnitStateTest {
+
+	@Test
+	public void testImageExtractionBasic()
+	{
+		final String name = "test";
+		PredictiveUnit pu = PredictiveUnit.newBuilder().setName(name).setEndpoint(Endpoint.newBuilder().setServiceHost("host")).build();
+		Map<String,Container> containersMap = new HashMap<String,Container>();
+		containersMap.put(name, Container.newBuilder().setImage("myimage:0.1").build());
+		PredictiveUnitState pus = new PredictiveUnitState(pu, containersMap);
+		Assert.assertEquals("myimage", pus.imageName);
+		Assert.assertEquals("0.1", pus.imageVersion);
+	}
+
+	@Test
+	public void testImageExtractionWithPort()
+	{
+		final String name = "test";
+		PredictiveUnit pu = PredictiveUnit.newBuilder().setName(name).setEndpoint(Endpoint.newBuilder().setServiceHost("host")).build();
+		Map<String,Container> containersMap = new HashMap<String,Container>();
+		containersMap.put(name, Container.newBuilder().setImage("myrep:1234/someorg/myimage:0.1").build());
+		PredictiveUnitState pus = new PredictiveUnitState(pu, containersMap);
+		Assert.assertEquals("myrep:1234/someorg/myimage", pus.imageName);
+		Assert.assertEquals("0.1", pus.imageVersion);
+	}
+
+}