From eaa72b7a2553361cde0f16d6d1eef815f558e0f5 Mon Sep 17 00:00:00 2001
From: ishaansehgal99 <ishaanforthewin@gmail.com>
Date: Fri, 12 Jul 2024 17:35:47 -0400
Subject: [PATCH] fix: Add adapter logs

---
 test/e2e/inference_with_adapters.go |  9 +++++++++
 test/e2e/preset_test.go             | 31 +++--------------------------
 test/e2e/utils/utils.go             | 25 +++++++++++++++++++++++
 3 files changed, 37 insertions(+), 28 deletions(-)

diff --git a/test/e2e/inference_with_adapters.go b/test/e2e/inference_with_adapters.go
index e07507635..5b3e8fc53 100644
--- a/test/e2e/inference_with_adapters.go
+++ b/test/e2e/inference_with_adapters.go
@@ -79,6 +79,15 @@ var _ = Describe("Workspace Preset", func() {
 		loadModelVersions()
 	})
 
+	AfterEach(func() {
+		if CurrentSpecReport().Failed() {
+			utils.PrintPodLogsOnFailure(namespaceName, "")     // The Preset Pod
+			utils.PrintPodLogsOnFailure("kaito-workspace", "") // The Kaito Workspace Pod
+			utils.PrintPodLogsOnFailure("gpu-provisioner", "") // The gpu-provisioner Pod
+			Fail("Fail threshold reached")
+		}
+	})
+
 	It("should create a falcon workspace with adapter", func() {
 		numOfNode := 1
 		workspaceObj := createCustomWorkspaceWithAdapter(numOfNode)
diff --git a/test/e2e/preset_test.go b/test/e2e/preset_test.go
index ff424cb59..8a129ad6b 100644
--- a/test/e2e/preset_test.go
+++ b/test/e2e/preset_test.go
@@ -533,31 +533,6 @@ func deleteWorkspace(workspaceObj *kaitov1alpha1.Workspace) error {
 	return nil
 }
 
-func printPodLogsOnFailure(namespace, labelSelector string) {
-	coreClient, err := utils.GetK8sConfig()
-	if err != nil {
-		log.Printf("Failed to create core client: %v", err)
-	}
-	pods, err := coreClient.CoreV1().Pods(namespace).List(ctx, metav1.ListOptions{
-		LabelSelector: labelSelector,
-	})
-	if err != nil {
-		log.Printf("Failed to list pods: %v", err)
-		return
-	}
-
-	for _, pod := range pods.Items {
-		for _, container := range pod.Spec.Containers {
-			logs, err := utils.GetPodLogs(coreClient, namespace, pod.Name, container.Name)
-			if err != nil {
-				log.Printf("Failed to get logs from pod %s, container %s: %v", pod.Name, container.Name, err)
-			} else {
-				fmt.Printf("Logs from pod %s, container %s:\n%s\n", pod.Name, container.Name, string(logs))
-			}
-		}
-	}
-}
-
 var runLlama13B bool
 var aiModelsRegistry string
 var aiModelsRegistrySecret string
@@ -574,9 +549,9 @@ var _ = Describe("Workspace Preset", func() {
 
 	AfterEach(func() {
 		if CurrentSpecReport().Failed() {
-			printPodLogsOnFailure(namespaceName, "")     // The Preset Pod
-			printPodLogsOnFailure("kaito-workspace", "") // The Kaito Workspace Pod
-			printPodLogsOnFailure("gpu-provisioner", "") // The gpu-provisioner Pod
+			utils.PrintPodLogsOnFailure(namespaceName, "")     // The Preset Pod
+			utils.PrintPodLogsOnFailure("kaito-workspace", "") // The Kaito Workspace Pod
+			utils.PrintPodLogsOnFailure("gpu-provisioner", "") // The gpu-provisioner Pod
 			Fail("Fail threshold reached")
 		}
 	})
diff --git a/test/e2e/utils/utils.go b/test/e2e/utils/utils.go
index 4aab4c0c2..b479b5df1 100644
--- a/test/e2e/utils/utils.go
+++ b/test/e2e/utils/utils.go
@@ -132,6 +132,31 @@ func GetPodLogs(coreClient *kubernetes.Clientset, namespace, podName, containerN
 	return buf.String(), nil
 }
 
+func PrintPodLogsOnFailure(namespace, labelSelector string) {
+	coreClient, err := GetK8sConfig()
+	if err != nil {
+		log.Printf("Failed to create core client: %v", err)
+	}
+	pods, err := coreClient.CoreV1().Pods(namespace).List(context.Background(), metav1.ListOptions{
+		LabelSelector: labelSelector,
+	})
+	if err != nil {
+		log.Printf("Failed to list pods: %v", err)
+		return
+	}
+
+	for _, pod := range pods.Items {
+		for _, container := range pod.Spec.Containers {
+			logs, err := GetPodLogs(coreClient, namespace, pod.Name, container.Name)
+			if err != nil {
+				log.Printf("Failed to get logs from pod %s, container %s: %v", pod.Name, container.Name, err)
+			} else {
+				fmt.Printf("Logs from pod %s, container %s:\n%s\n", pod.Name, container.Name, string(logs))
+			}
+		}
+	}
+}
+
 func CopySecret(original *corev1.Secret, targetNamespace string) *corev1.Secret {
 	newSecret := &corev1.Secret{
 		ObjectMeta: metav1.ObjectMeta{