Add tests

Author: quoctruong

packages/k8s/src/hooks/prepare-job.ts

@@ -144,7 +144,7 @@ export async function prepareJob(
   generateResponseFile(responseFile, args, createdPod, isAlpine)
 }
 
-function processServiceContainers(
+export function processServiceContainers(
   services?: ServiceContainerInfo[],
   container?: k8s.V1Container,
   extension?: k8s.V1PodTemplateSpec
@@ -157,7 +157,7 @@ function processServiceContainers(
     core.debug(`Adding service '${service.image}' to pod definition`)
     return createContainerSpec(
       service,
-      generateContainerName(service.image),
+      service.name,
       false,
       extension,
       service.createOptions
@@ -188,6 +188,7 @@ function processServiceContainers(
         container.resources.limits &&
         container.resources.limits['google.com/tpu']
       ) {
+        core.debug('removing tpu from main container resource limits')
         delete container.resources.limits['google.com/tpu']
       }
     }
@@ -399,11 +400,20 @@ export function createContainerSpec(
   }
 
   podContainer.env = []
-  for (const [key, value] of Object.entries(
-    container['environmentVariables']
-  )) {
-    if (value && key !== 'HOME') {
-      podContainer.env.push({ name: key, value: value as string })
+  if (container['environmentVariables']) {
+    for (const [key, value] of Object.entries(
+      container['environmentVariables']
+    )) {
+      if (value && key !== 'HOME') {
+        podContainer.env.push({ name: key, value: value as string })
+      }
+    }
+
+    if (!('CI' in container['environmentVariables'])) {
+      podContainer.env.push({
+        name: 'CI',
+        value: 'true'
+      })
     }
   }
 
@@ -412,13 +422,6 @@ export function createContainerSpec(
     value: 'true'
   })
 
-  if (!('CI' in container['environmentVariables'])) {
-    podContainer.env.push({
-      name: 'CI',
-      value: 'true'
-    })
-  }
-
   podContainer.volumeMounts = containerVolumes(
     container.userMountVolumes,
     jobContainer

packages/k8s/tests/prepare-job-test.ts

@@ -1,18 +1,20 @@
 import * as fs from 'fs'
 import * as path from 'path'
 import { cleanupJob } from '../src/hooks'
-import { createContainerSpec, prepareJob } from '../src/hooks/prepare-job'
+import {
+  createContainerSpec,
+  prepareJob,
+  processServiceContainers
+} from '../src/hooks/prepare-job'
 import { TestHelper } from './test-setup'
 import {
   ENV_HOOK_TEMPLATE_PATH,
   ENV_NUMBER_OF_HOSTS,
   ENV_USE_KUBE_SCHEDULER,
-  generateContainerName,
-  readExtensionFromFile
+  generateContainerName
 } from '../src/k8s/utils'
-import { getEvents, getPodByName } from '../src/k8s'
+import { getPodByName } from '../src/k8s'
 import { V1Container } from '@kubernetes/client-node'
-import * as yaml from 'js-yaml'
 import { JOB_CONTAINER_NAME } from '../src/hooks/constants'
 
 jest.useRealTimers()
@@ -324,3 +326,62 @@ describe('Prepare job', () => {
     }
   )
 })
+
+describe('processServiceContainers', () => {
+  it('generate names for service containers', () => {
+    expect(
+      processServiceContainers(
+        [
+          {
+            image: 'gcr.io/server'
+          },
+          {
+            image: 'gcr.io/server'
+          }
+        ],
+        {
+          name: 'nginx',
+          image: 'nginx:latest',
+          imagePullPolicy: 'IfNotPresent'
+        } as V1Container
+      )
+    ).toEqual(
+      expect.arrayContaining([
+        expect.objectContaining({ name: 'server' }),
+        expect.objectContaining({ name: 'server-1' })
+      ])
+    )
+  })
+
+  it('generate TPU request for service containers', () => {
+    expect(
+      processServiceContainers(
+        [
+          {
+            image: 'gcr.io/server',
+            createOptions: '--tpu=4'
+          }
+        ],
+        {
+          name: 'nginx',
+          image: 'nginx:latest',
+          imagePullPolicy: 'IfNotPresent'
+        } as V1Container
+      )
+    ).toEqual(
+      expect.arrayContaining([
+        expect.objectContaining({
+          name: 'server',
+          resources: {
+            limits: {
+              'google.com/tpu': '4'
+            },
+            requests: {
+              'google.com/tpu': '4'
+            }
+          }
+        })
+      ])
+    )
+  })
+})