Adjust 'go.embed' to accept python files

Author: ChughShilpa

tests/kfto/core/kfto_training_test.go

@@ -46,7 +46,7 @@ func runKFTOPyTorchJob(t *testing.T, image string, gpuLabel string, numGpus int)
 
 	// Create a ConfigMap with training script
 	configData := map[string][]byte{
-		"hf_llm_training.py": ReadFileExt(test, "hf_llm_training.py"),
+		"hf_llm_training.py": ReadFile(test, "hf_llm_training.py"),
 	}
 	config := CreateConfigMap(test, namespace, configData)
 

tests/kfto/core/support.go

@@ -19,10 +19,8 @@ package core
 import (
 	"embed"
 	"fmt"
-	"os"
 	"time"
 
-	"github.com/onsi/gomega"
 	. "github.com/onsi/gomega"
 	. "github.com/project-codeflare/codeflare-common/support"
 
@@ -34,6 +32,7 @@ import (
 )
 
 //go:embed *.json
+//go:embed *.py
 var files embed.FS
 
 func ReadFile(t Test, fileName string) []byte {
@@ -43,13 +42,6 @@ func ReadFile(t Test, fileName string) []byte {
 	return file
 }
 
-func ReadFileExt(t Test, fileName string) []byte {
-	t.T().Helper()
-	file, err := os.ReadFile(fileName)
-	t.Expect(err).NotTo(gomega.HaveOccurred())
-	return file
-}
-
 func PyTorchJob(t Test, namespace, name string) func(g Gomega) *kftov1.PyTorchJob {
 	return func(g Gomega) *kftov1.PyTorchJob {
 		job, err := t.Client().Kubeflow().KubeflowV1().PyTorchJobs(namespace).Get(t.Ctx(), name, metav1.GetOptions{})