github.com/NVIDIA/aistore@v1.3.23-0.20240517131212-7df6609be51d/ext/etl/comm_internal_test.go (about)

     1  // Package etl provides utilities to initialize and use transformation pods.
     2  /*
     3   * Copyright (c) 2018-2023, NVIDIA CORPORATION. All rights reserved.
     4   */
     5  package etl
     6  
     7  import (
     8  	"fmt"
     9  	"io"
    10  	"net/http"
    11  	"net/http/httptest"
    12  	"os"
    13  	"path/filepath"
    14  	"testing"
    15  	"time"
    16  
    17  	"github.com/NVIDIA/aistore/api/apc"
    18  	"github.com/NVIDIA/aistore/cmn"
    19  	"github.com/NVIDIA/aistore/cmn/cos"
    20  	"github.com/NVIDIA/aistore/core"
    21  	"github.com/NVIDIA/aistore/core/meta"
    22  	"github.com/NVIDIA/aistore/core/mock"
    23  	"github.com/NVIDIA/aistore/fs"
    24  	"github.com/NVIDIA/aistore/tools/cryptorand"
    25  	. "github.com/onsi/ginkgo/v2"
    26  	. "github.com/onsi/gomega"
    27  	corev1 "k8s.io/api/core/v1"
    28  )
    29  
    30  func TestETLTransform(t *testing.T) {
    31  	if testing.Short() {
    32  		t.Skipf("skipping %s in short mode", t.Name())
    33  	}
    34  	RegisterFailHandler(Fail)
    35  	RunSpecs(t, t.Name())
    36  }
    37  
    38  var _ = Describe("CommunicatorTest", func() {
    39  	var (
    40  		tmpDir            string
    41  		comm              Communicator
    42  		transformerServer *httptest.Server
    43  		targetServer      *httptest.Server
    44  		proxyServer       *httptest.Server
    45  
    46  		dataSize      = int64(cos.MiB * 50)
    47  		transformData = make([]byte, dataSize)
    48  
    49  		bck        = cmn.Bck{Name: "commBck", Provider: apc.AIS, Ns: cmn.NsGlobal}
    50  		objName    = "commObj"
    51  		clusterBck = meta.NewBck(
    52  			bck.Name, bck.Provider, bck.Ns,
    53  			&cmn.Bprops{Cksum: cmn.CksumConf{Type: cos.ChecksumXXHash}},
    54  		)
    55  		bmdMock = mock.NewBaseBownerMock(clusterBck)
    56  	)
    57  
    58  	BeforeEach(func() {
    59  		// Generate random data.
    60  		_, err := cryptorand.Read(transformData)
    61  		Expect(err).NotTo(HaveOccurred())
    62  
    63  		// Initialize the filesystem and directories.
    64  		tmpDir, err = os.MkdirTemp("", "")
    65  		Expect(err).NotTo(HaveOccurred())
    66  
    67  		mpath := filepath.Join(tmpDir, "mpath")
    68  		err = cos.CreateDir(mpath)
    69  		Expect(err).NotTo(HaveOccurred())
    70  		fs.TestNew(nil)
    71  		_, err = fs.Add(mpath, "daeID")
    72  		Expect(err).NotTo(HaveOccurred())
    73  
    74  		_ = mock.NewTarget(bmdMock)
    75  		// cluster.InitLomLocker(tMock)
    76  
    77  		// Create an object.
    78  		lom := &core.LOM{ObjName: objName}
    79  		err = lom.InitBck(clusterBck.Bucket())
    80  		Expect(err).NotTo(HaveOccurred())
    81  		err = createRandomFile(lom.FQN, dataSize)
    82  		Expect(err).NotTo(HaveOccurred())
    83  		lom.SetAtimeUnix(time.Now().UnixNano())
    84  		lom.SetSize(dataSize)
    85  		err = lom.Persist()
    86  		Expect(err).NotTo(HaveOccurred())
    87  
    88  		// Initialize the HTTP servers.
    89  		transformerServer = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
    90  			_, err := w.Write(transformData)
    91  			Expect(err).NotTo(HaveOccurred())
    92  		}))
    93  		targetServer = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    94  			err := comm.InlineTransform(w, r, clusterBck, objName)
    95  			Expect(err).NotTo(HaveOccurred())
    96  		}))
    97  		proxyServer = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    98  			http.Redirect(w, r, targetServer.URL, http.StatusMovedPermanently)
    99  		}))
   100  	})
   101  
   102  	AfterEach(func() {
   103  		_ = os.RemoveAll(tmpDir)
   104  		proxyServer.Close()
   105  		transformerServer.Close()
   106  		targetServer.Close()
   107  	})
   108  
   109  	tests := []string{
   110  		Hpush,
   111  		Hpull,
   112  		Hrev,
   113  	}
   114  
   115  	for _, commType := range tests {
   116  		It("should perform transformation "+commType, func() {
   117  			pod := &corev1.Pod{}
   118  			pod.SetName("somename")
   119  
   120  			xctn := mock.NewXact(apc.ActETLInline)
   121  			boot := &etlBootstrapper{
   122  				msg: InitSpecMsg{
   123  					InitMsgBase: InitMsgBase{
   124  						CommTypeX: commType,
   125  					},
   126  				},
   127  				pod:  pod,
   128  				uri:  transformerServer.URL,
   129  				xctn: xctn,
   130  			}
   131  			comm = newCommunicator(nil, boot)
   132  
   133  			resp, err := http.Get(proxyServer.URL)
   134  			Expect(err).NotTo(HaveOccurred())
   135  			defer resp.Body.Close()
   136  
   137  			b, err := io.ReadAll(resp.Body)
   138  			Expect(err).NotTo(HaveOccurred())
   139  			Expect(len(b)).To(Equal(len(transformData)))
   140  			Expect(b).To(Equal(transformData))
   141  		})
   142  	}
   143  })
   144  
   145  // Creates a file with random content.
   146  func createRandomFile(fileName string, size int64) error {
   147  	b := make([]byte, size)
   148  	if _, err := cryptorand.Read(b); err != nil {
   149  		return err
   150  	}
   151  	f, err := cos.CreateFile(fileName)
   152  	if err != nil {
   153  		return err
   154  	}
   155  	defer f.Close()
   156  	n, err := f.Write(b)
   157  	if err != nil {
   158  		return err
   159  	}
   160  	if int64(n) != size {
   161  		return fmt.Errorf("could not write %d bytes, wrote %d bytes", size, n)
   162  	}
   163  	return nil
   164  }