test(computing-unit): cover the pod lookup, creation and deletion paths

computing-unit-managing-service/src/test/scala/org/apache/texera/service/util/KubernetesClientSpec.scala

@@ -26,17 +26,29 @@ import io.fabric8.kubernetes.api.model.metrics.v1beta1.{
   PodMetricsList,
   PodMetricsListBuilder
 }
-import io.fabric8.kubernetes.api.model.{Pod, PodBuilder, PodList, PodListBuilder, Quantity}
+import io.fabric8.kubernetes.api.model.{
+  ContainerBuilder,
+  Pod,
+  PodBuilder,
+  PodList,
+  PodListBuilder,
+  Quantity,
+  ResourceRequirementsBuilder
+}
 import io.fabric8.kubernetes.client.dsl.{
   MetricAPIGroupDSL,
   MixedOperation,
+  NamespaceableResource,
   NonNamespaceOperation,
   PodMetricOperation,
-  PodResource
+  PodResource,
+  Resource
 }
 import io.fabric8.kubernetes.client.{KubernetesClient => Fabric8Client}
 import org.apache.texera.common.config.KubernetesConfig
-import org.mockito.Mockito.{mock, when}
+import org.mockito.ArgumentCaptor
+import org.mockito.ArgumentMatchers.any
+import org.mockito.Mockito.{mock, times, verify, when}
 import org.scalatest.flatspec.AnyFlatSpec
 import org.scalatest.matchers.should.Matchers
 
@@ -160,4 +172,149 @@ class KubernetesClientSpec extends AnyFlatSpec with Matchers {
     k8s.getPodMetrics(1) shouldBe Map("cpu" -> "250m", "memory" -> "128Mi")
     k8s.getPodMetrics(999) shouldBe empty
   }
+  // ── single-pod lookups, creation and deletion ──
+  // These reach the rest of the fluent chain: withName(...).get() for the lookups,
+  // resource(pod).inNamespace(...).create() for creation, and .delete() for removal. Everything
+  // is driven through the constructor seam, so no cluster is involved.
+
+  /** Extends the namespace stub with the by-name pod operations `withName(...)` returns. */
+  private def clientWithNamedPod(podName: String, found: Pod): (Fabric8Client, PodResource) = {
+    val client = stubbedClient(Seq.empty, Seq.empty)
+    val podsInNamespace = client.pods().inNamespace(namespace)
+    val podResource = mock(classOf[PodResource])
+    when(podsInNamespace.withName(podName)).thenReturn(podResource)
+    when(podResource.get()).thenReturn(found)
+    (client, podResource)
+  }
+
+  /** A pod carrying one container whose resource limits are set. */
+  private def podWithLimits(cuid: Int, limits: Map[String, String]): Pod =
+    new PodBuilder()
+      .withNewMetadata()
+      .withName(KubernetesClient.generatePodName(cuid))
+      .endMetadata()
+      .withNewSpec()
+      .addToContainers(
+        new ContainerBuilder()
+          .withName("main")
+          .withResources(
+            new ResourceRequirementsBuilder()
+              .withLimits(limits.map { case (k, v) => k -> new Quantity(v) }.asJava)
+              .build()
+          )
+          .build()
+      )
+      .endSpec()
+      .build()
+
+  "generatePodURI" should "address the pod through its headless service inside the namespace" in {
+    // The URI is how a computing unit is reached once it is up, so every segment matters: a pod
+    // name alone, or the wrong namespace, resolves to nothing.
+    val uri = KubernetesClient.generatePodURI(7)
+    uri should startWith(KubernetesClient.generatePodName(7) + ".")
+    uri should include(s".${KubernetesConfig.computeUnitServiceName}.$namespace.svc.cluster.local:")
+    uri should endWith(s":${KubernetesConfig.computeUnitPortNumber}")
+  }
+
+  "getPodByName" should "wrap a found pod and report a missing one as None" in {
+    // fabric8 returns null rather than throwing for an absent pod, so the Option() wrapper is the
+    // only thing standing between a caller and an NPE.
+    val name = KubernetesClient.generatePodName(1)
+    val (found, _) = clientWithNamedPod(name, pod(1, "Running"))
+    val (absent, _) = clientWithNamedPod(name, null)
+
+    new KubernetesClient(found).getPodByName(name).map(_.getMetadata.getName) shouldBe Some(name)
+    new KubernetesClient(absent).getPodByName(name) shouldBe None
+  }
+
+  "podExists" should "follow the by-name lookup in both directions" in {
+    val name = KubernetesClient.generatePodName(2)
+    new KubernetesClient(clientWithNamedPod(name, pod(2, "Running"))._1).podExists(2) shouldBe true
+    new KubernetesClient(clientWithNamedPod(name, null)._1).podExists(2) shouldBe false
+  }
+
+  "getPodLimits" should "read the first container's limits and fall back to an empty map" in {
+    val name = KubernetesClient.generatePodName(3)
+    val withLimits =
+      clientWithNamedPod(name, podWithLimits(3, Map("cpu" -> "2", "memory" -> "4Gi")))._1
+    val missing = clientWithNamedPod(name, null)._1
+
+    new KubernetesClient(withLimits).getPodLimits(3) shouldBe Map("cpu" -> "2", "memory" -> "4Gi")
+    new KubernetesClient(missing).getPodLimits(3) shouldBe empty
+  }
+
+  "createPod" should "refuse to overwrite a pod that already exists" in {
+    // Creating over a live unit would silently detach the running one from its owner.
+    val name = KubernetesClient.generatePodName(4)
+    val k8s = new KubernetesClient(clientWithNamedPod(name, pod(4, "Running"))._1)
+
+    val thrown = intercept[Exception] {
+      k8s.createPod(4, "1", "2Gi", "0", Map.empty)
+    }
+    thrown.getMessage should include("already exists")
+  }
+
+  it should "build the pod from the requested limits and env, and create it in the namespace" in {
+    val name = KubernetesClient.generatePodName(5)
+    val (client, _) = clientWithNamedPod(name, null)
+    val namespaceable = mock(classOf[NamespaceableResource[Pod]])
+    val resource = mock(classOf[Resource[Pod]])
+    val captor = ArgumentCaptor.forClass(classOf[Pod])
+    when(client.resource(any(classOf[Pod]))).thenReturn(namespaceable)
+    when(namespaceable.inNamespace(namespace)).thenReturn(resource)
+    // create()'s return value is not asserted; the pod is inspected through the captor below.
+    when(resource.create()).thenReturn(null)
+
+    new KubernetesClient(client).createPod(5, "2", "4Gi", "1", Map("UID" -> 9, "MODE" -> "batch"))
+
+    verify(client).resource(captor.capture())
+    val built = captor.getValue
+    built.getSpec.getHostname shouldBe name
+    built.getSpec.getSubdomain shouldBe KubernetesConfig.computeUnitServiceName
+    val container = built.getSpec.getContainers.asScala.head
+    val limits = container.getResources.getLimits.asScala.map { case (k, v) => k -> v.toString }
+    limits("cpu") shouldBe "2"
+    limits("memory") shouldBe "4Gi"
+    // Env values arrive as Any and reach the container as strings.
+    container.getEnv.asScala.map(e => e.getName -> e.getValue).toMap shouldBe
+      Map("UID" -> "9", "MODE" -> "batch")
+  }
+
+  it should "mount a shared-memory volume only when a size is asked for" in {
+    // /dev/shm defaults to 64Mi in Kubernetes, which is too small for the Python workers, so the
+    // volume is the fix — but it must not appear when no size was requested.
+    def build(shm: Option[String]): Pod = {
+      val name = KubernetesClient.generatePodName(6)
+      val (client, _) = clientWithNamedPod(name, null)
+      val namespaceable = mock(classOf[NamespaceableResource[Pod]])
+      val resource = mock(classOf[Resource[Pod]])
+      val captor = ArgumentCaptor.forClass(classOf[Pod])
+      when(client.resource(any(classOf[Pod]))).thenReturn(namespaceable)
+      when(namespaceable.inNamespace(namespace)).thenReturn(resource)
+      // create()'s return value is not asserted; the pod is inspected through the captor below.
+      when(resource.create()).thenReturn(null)
+      new KubernetesClient(client).createPod(6, "1", "2Gi", "0", Map.empty, shm)
+      verify(client).resource(captor.capture())
+      captor.getValue
+    }
+
+    val withShm = build(Some("1Gi"))
+    withShm.getSpec.getVolumes.asScala.map(_.getName) should contain("dshm")
+    withShm.getSpec.getVolumes.asScala
+      .find(_.getName == "dshm")
+      .flatMap(v => Option(v.getEmptyDir))
+      .map(_.getSizeLimit.toString) shouldBe Some("1Gi")
+
+    Option(build(None).getSpec.getVolumes).map(_.asScala.map(_.getName)).getOrElse(Nil) should
+      not contain "dshm"
+  }
+
+  "deletePod" should "delete the pod for the cuid inside the namespace" in {
+    val name = KubernetesClient.generatePodName(8)
+    val (client, podResource) = clientWithNamedPod(name, pod(8, "Running"))
+
+    new KubernetesClient(client).deletePod(8)
+
+    verify(podResource, times(1)).delete()
+  }
 }