From ca4ea17d58bec3a2a3b8ea28347f76f575fbf04b Mon Sep 17 00:00:00 2001
From: Nikola Jokic <jokicnikola07@gmail.com>
Date: Wed, 19 Jun 2024 11:49:43 +0200
Subject: [PATCH] Skip writing extension containers in output context file
 (#154)

---
 packages/k8s/src/hooks/prepare-job.ts  | 38 ++++++++++++++------------
 packages/k8s/tests/prepare-job-test.ts | 20 ++++++++++++++
 2 files changed, 41 insertions(+), 17 deletions(-)

diff --git a/packages/k8s/src/hooks/prepare-job.ts b/packages/k8s/src/hooks/prepare-job.ts
index 7f3d81b..89202dc 100644
--- a/packages/k8s/src/hooks/prepare-job.ts
+++ b/packages/k8s/src/hooks/prepare-job.ts
@@ -119,11 +119,12 @@ export async function prepareJob(
     throw new Error(`failed to determine if the pod is alpine: ${message}`)
   }
   core.debug(`Setting isAlpine to ${isAlpine}`)
-  generateResponseFile(responseFile, createdPod, isAlpine)
+  generateResponseFile(responseFile, args, createdPod, isAlpine)
 }
 
 function generateResponseFile(
   responseFile: string,
+  args: PrepareJobArgs,
   appPod: k8s.V1Pod,
   isAlpine
 ): void {
@@ -156,24 +157,27 @@ function generateResponseFile(
     }
   }
 
-  const serviceContainers = appPod.spec?.containers.filter(
-    c => c.name !== JOB_CONTAINER_NAME
-  )
-  if (serviceContainers?.length) {
-    response.context['services'] = serviceContainers.map(c => {
-      const ctxPorts: ContextPorts = {}
-      if (c.ports?.length) {
-        for (const port of c.ports) {
-          ctxPorts[port.containerPort] = port.hostPort
-        }
-      }
+  if (args.services?.length) {
+    const serviceContainerNames =
+      args.services?.map(s => generateContainerName(s.image)) || []
 
-      return {
-        image: c.image,
-        ports: ctxPorts
-      }
-    })
+    response.context['services'] = appPod?.spec?.containers
+      ?.filter(c => serviceContainerNames.includes(c.name))
+      .map(c => {
+        const ctxPorts: ContextPorts = {}
+        if (c.ports?.length) {
+          for (const port of c.ports) {
+            ctxPorts[port.containerPort] = port.hostPort
+          }
+        }
+
+        return {
+          image: c.image,
+          ports: ctxPorts
+        }
+      })
   }
+
   writeToResponseFile(responseFile, JSON.stringify(response))
 }
 
diff --git a/packages/k8s/tests/prepare-job-test.ts b/packages/k8s/tests/prepare-job-test.ts
index 71336df..671294c 100644
--- a/packages/k8s/tests/prepare-job-test.ts
+++ b/packages/k8s/tests/prepare-job-test.ts
@@ -147,6 +147,26 @@ describe('Prepare job', () => {
     expect(got.spec?.containers[2].args).toEqual(['-c', 'sleep 60'])
   })
 
+  it('should put only job and services in output context file', async () => {
+    process.env[ENV_HOOK_TEMPLATE_PATH] = path.join(
+      __dirname,
+      '../../../examples/extension.yaml'
+    )
+
+    await expect(
+      prepareJob(prepareJobData.args, prepareJobOutputFilePath)
+    ).resolves.not.toThrow()
+
+    const content = JSON.parse(
+      fs.readFileSync(prepareJobOutputFilePath).toString()
+    )
+
+    expect(content.state.jobPod).toBeTruthy()
+    expect(content.context.container).toBeTruthy()
+    expect(content.context.services).toBeTruthy()
+    expect(content.context.services.length).toBe(1)
+  })
+
   it('should not throw exception using kube scheduler', async () => {
     // only for ReadWriteMany volumes or single node cluster
     process.env[ENV_USE_KUBE_SCHEDULER] = 'true'