From 3fe2938765a678e67b71e4aca32901a6fe31577c Mon Sep 17 00:00:00 2001
From: Santanu Sinha <santanu.sinha@gmail.com>
Date: Sat, 13 Sep 2025 19:22:38 +0530
Subject: [PATCH] Pulling container before calling create in test

---
 .../drove/executor/utils/DockerUtilsTest.java     | 15 ++++++++++++++-
 1 file changed, 14 insertions(+), 1 deletion(-)

diff --git a/drove-executor/src/test/java/com/phonepe/drove/executor/utils/DockerUtilsTest.java b/drove-executor/src/test/java/com/phonepe/drove/executor/utils/DockerUtilsTest.java
index e45a2d5b..e0e4c1b3 100644
--- a/drove-executor/src/test/java/com/phonepe/drove/executor/utils/DockerUtilsTest.java
+++ b/drove-executor/src/test/java/com/phonepe/drove/executor/utils/DockerUtilsTest.java
@@ -38,6 +38,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.function.Function;
 import java.util.stream.Stream;
 
 import static com.phonepe.drove.executor.ExecutorTestingUtils.DOCKER_CLIENT;
@@ -65,7 +66,7 @@ void testContainerCreate(
         when(resourceManager.currentState())
                 .thenReturn(new ResourceInfo(null, null,
                                              new PhysicalLayout(Map.of(0, Set.of(0, 1, 2, 3)), Map.of(0, 1024L))));
-        val httpCaller = mock(HttpCaller.class);
+        pullImage(spec);
         val containerId = DockerUtils.createContainer(
                 resourceConfig,
                 DOCKER_CLIENT,
@@ -91,6 +92,7 @@ void testContainerRun() {
         val httpCaller = mock(HttpCaller.class);
         val spec = ExecutorTestingUtils.testTaskInstanceSpec()
                 .withEnv(Map.of("ITERATIONS", "10000", "LOOP_SLEEP", "5"));
+        pullImage(spec);
         val containerId = DockerUtils.createContainer(
                 ResourceConfig.DEFAULT,
                 DOCKER_CLIENT,
@@ -114,6 +116,17 @@ void testContainerRun() {
         }
     }
 
+    private static void pullImage(TaskInstanceSpec spec) {
+        try {
+            val imageId = DockerUtils.pullImage(DOCKER_CLIENT, null, spec, Function.identity());
+            assertNotNull(DOCKER_CLIENT.inspectImageCmd(imageId).exec());
+        }
+        catch (InterruptedException e) {
+            Thread.currentThread().interrupt();
+            Assertions.fail("Pull Interrupted", e);
+        }
+    }
+
     private static Stream<Arguments> generateResourceConfigs() {
         val spec = testTaskInstanceSpec();
         val defaultConfig = ResourceConfig.DEFAULT;