tools/gcs/gcs_test.go

package gcs

import (
	"context"
	"fmt"
	"os"
	"path/filepath"
	"reflect"
	"strings"
	"testing"

	"github.com/fsouza/fake-gcs-server/fakestorage"
)

func TestGetModulePath(t *testing.T) {
	tmpDir, err := os.MkdirTemp("", "test-getmodulepath")
	if err != nil {
		t.Fatalf("Failed to create temp dir: %v", err)
	}
	defer os.RemoveAll(tmpDir)


	tests := []struct {
		desc       string
		localDir   string
		moduleName string
		wantErr    bool
		wantPath   string
	}{
		{
			desc:       "Valid module name",
			localDir:   tmpDir,
			moduleName: "module.ko",
			wantErr:    false,
			wantPath:   filepath.Join(tmpDir, "module.ko"),
		},
		{
			desc:       "Valid module name with subdirectory",
			localDir:   tmpDir,
			moduleName: "subdir/module.ko",
			wantErr:    false,
			wantPath:   filepath.Join(tmpDir, "subdir/module.ko"),
		},
		{
			desc:       "Path traversal",
			localDir:   tmpDir,
			moduleName: "../../../module.ko",
			wantErr:    true,
		},
		{
			desc:       "Empty module name",
			localDir:   tmpDir,
			moduleName: "",
			wantErr:    true,
		},
	}

	for _, test := range tests {
		t.Run(test.desc, func(t *testing.T) {
			gotPath, err := getModulePath(test.localDir, test.moduleName)
			if gotErr := err != nil; gotErr != test.wantErr {
				t.Errorf("getModulePath(%q, %q): expected error: %t, got error: %v", test.localDir, test.moduleName, test.wantErr, err)
			}

			if !test.wantErr {
				// Normalize paths for comparison
				wantAbsPath, _ := filepath.Abs(test.wantPath)
				gotAbsPath, _ := filepath.Abs(gotPath)
				if gotAbsPath != wantAbsPath {
					t.Errorf("getModulePath(%q, %q): expected path: %q, got path: %q", test.localDir, test.moduleName, wantAbsPath, gotAbsPath)
				}

				// Extra check to ensure the path is within the base directory
				absLocalDir, _ := filepath.Abs(test.localDir)
				if !strings.HasPrefix(gotAbsPath, absLocalDir) {
					t.Errorf("getModulePath(%q, %q): resulting path %q is not under %q", test.localDir, test.moduleName, gotAbsPath, absLocalDir)
				}
			}
		})
	}
}

func TestGetBucketAndDir(t *testing.T) {
	tests := []struct {
		desc       string
		gsUri      string
		wantErr    bool
		wantBucket string
		wantDir    string
	}{
		{
			desc:    "Invalid path syntax",
			gsUri:   "https://storage.mtls.cloud.google.com/gcs/my-bucket/folder",
			wantErr: true,
		},
		{
			desc:    "Missing gcs path",
			gsUri:   "",
			wantErr: true,
		},
		{
			desc:    "Missing bucket name",
			gsUri:   "gs://",
			wantErr: true,
		},
		{
			desc:    "Missing folder ending with /",
			gsUri:   "gs://bucket-name/",
			wantErr: true,
		},
		{
			desc:    "Missing folder",
			gsUri:   "gs://bucket-name",
			wantErr: true,
		},
		{
			desc:       "Valid gcs path",
			gsUri:      "gs://my-bucket/path/to/modules",
			wantBucket: "my-bucket",
			wantDir:    "path/to/modules/",
		},
		{
			desc:       "Valid gcs path ending with /",
			gsUri:      "gs://my-bucket/path/to/modules/",
			wantBucket: "my-bucket",
			wantDir:    "path/to/modules/",
		},
	}
	for _, test := range tests {
		t.Run(test.desc, func(t *testing.T) {
			bktName, dir, err := getBucketNameAndDir(test.gsUri)
			if gotErr := err != nil; gotErr != test.wantErr {
				t.Errorf("TestGetBucketAndDir(%s): Error: %s\n expected error: %t\n got error: %t", test.desc, err, test.wantErr, gotErr)
			}
			if bktName != test.wantBucket {
				t.Errorf("TestGetBucketAndDir(%s): expected bucket:%s\t got bucket:%s", test.desc, test.wantBucket, bktName)
			}
			if dir != test.wantDir {
				t.Errorf("TestGetBucketAndDir(%s): expected directory:%s\t got directory:%s", test.desc, test.wantDir, dir)
			}
		})
	}
}

func TestInit(t *testing.T) {
	ctx := context.Background()
	tests := []struct {
		desc       string
		gsUri      string
		server     *fakestorage.Server
		wantErr    bool
		wantBucket string
		wantPath   string
	}{
		{
			desc:  "Bucket does not exist",
			gsUri: "gs://unknown/path/to/module",
			server: fakestorage.NewServer(
				[]fakestorage.Object{
					{
						ObjectAttrs: fakestorage.ObjectAttrs{
							BucketName: "my-bucket",
							Name:       "path/to/module/module1.ko",
						},
					},
				}),
			wantErr: true,
		},
		{
			desc:  "Successfully sets up the gcs config",
			gsUri: "gs://my-bucket/path/to/module",
			server: fakestorage.NewServer(
				[]fakestorage.Object{
					{
						ObjectAttrs: fakestorage.ObjectAttrs{
							BucketName: "my-bucket",
							Name:       "path/to/module/module1.ko",
						},
					},
				}),
			wantBucket: "my-bucket",
			wantPath:   "path/to/module/",
		},
	}
	for _, test := range tests {
		t.Run(test.desc, func(t *testing.T) {
			defer test.server.Stop()
			var cfg GCSConfig
			err := cfg.Init(ctx, test.gsUri, test.server.Client())
			if gotErr := err != nil; gotErr != test.wantErr {
				t.Errorf("TestInit(%s): Error: %v\n expected error: %t\n got error: %t", test.desc, err, test.wantErr, gotErr)
			}
			bucket := cfg.bucketName()
			if bucket != test.wantBucket {
				t.Errorf("TestListModules(%s): expected bucket:%s\t got bucket:%s", test.desc, test.wantBucket, bucket)
			}
			if cfg.path != test.wantPath {
				t.Errorf("TestListModules(%s): expected path:%s\t got path:%s", test.desc, test.wantPath, cfg.path)
			}

		})
	}
}

func TestInitNonExistentPath(t *testing.T) {
	ctx := context.Background()
	server := fakestorage.NewServer([]fakestorage.Object{})
	server.CreateBucketWithOpts(fakestorage.CreateBucketOpts{Name: "my-bucket"})
	defer server.Stop()
	var cfg GCSConfig
	err := cfg.Init(ctx, "gs://my-bucket/path/to/module", server.Client())
	if gotErr := err != nil; gotErr != true {
		t.Errorf("TestInitNonExistentPath: Error: %v\n expected error: true\n got error: %t", err, gotErr)
	}
}

func TestListModules(t *testing.T) {
	ctx := context.Background()
	tests := []struct {
		desc        string
		gcsPath     string
		server      *fakestorage.Server
		recursive   bool
		wantErr     bool
		wantModules []string
	}{
		{
			desc:    "Successfully retrieves modules",
			gcsPath: "path/to/module/",
			server: fakestorage.NewServer(
				[]fakestorage.Object{
					{
						ObjectAttrs: fakestorage.ObjectAttrs{
							BucketName: "my-bucket",
							Name:       "path/to/module/module1.ko",
						},
					},
					{
						ObjectAttrs: fakestorage.ObjectAttrs{
							BucketName: "my-bucket",
							Name:       "path/to/module/module2.ko",
						},
					},
					{
						ObjectAttrs: fakestorage.ObjectAttrs{
							BucketName: "my-bucket",
							Name:       "path/to/module/text.txt",
						},
					},
				}),
			wantModules: []string{"module1.ko", "module2.ko"},
		}, {
			desc:    "Does not return modules in subdirectory - recursive not set",
			gcsPath: "path/to/module/",
			server: fakestorage.NewServer(
				[]fakestorage.Object{
					{
						ObjectAttrs: fakestorage.ObjectAttrs{
							BucketName: "my-bucket",
							Name:       "path/to/module/module1.ko",
						},
					},
					{
						ObjectAttrs: fakestorage.ObjectAttrs{
							BucketName: "my-bucket",
							Name:       "path/to/module/subdirectory/module2.ko",
						},
					},
				}),
			wantModules: []string{"module1.ko"},
		},
		{
			desc:    "Retrieves modules in subdirectory - recursive set",
			gcsPath: "path/to/module/",
			server: fakestorage.NewServer(
				[]fakestorage.Object{
					{
						ObjectAttrs: fakestorage.ObjectAttrs{
							BucketName: "my-bucket",
							Name:       "path/to/module/module1.ko",
						},
					},
					{
						ObjectAttrs: fakestorage.ObjectAttrs{
							BucketName: "my-bucket",
							Name:       "path/to/module/subdirectory/module2.ko",
						},
					},
				}),
			recursive:   true,
			wantModules: []string{"module1.ko", "subdirectory/module2.ko"},
		},
	}
	for _, test := range tests {
		t.Run(test.desc, func(t *testing.T) {
			defer test.server.Stop()
			cfg := &GCSConfig{bucket: test.server.Client().Bucket("my-bucket"), path: test.gcsPath}
			modules, err := cfg.ListModules(ctx, test.recursive)
			if gotErr := err != nil; gotErr != test.wantErr {
				t.Errorf("TestListModules(%s): Error: %s\n expected error: %t\n got error: %t", test.desc, err, test.wantErr, gotErr)
			}
			if !reflect.DeepEqual(test.wantModules, modules) {
				t.Errorf("TestListModules(%s): expected modules: %s\t got modules: %s", test.desc, test.wantModules, modules)
			}
		})
	}
}

func TestDownloadKernelModuleSuccessful(t *testing.T) {
	ctx := context.Background()
	tmpDir, err := os.MkdirTemp("", "test-gcs")
	if err != nil {
		t.Fatalf("Failed to create temp dir: %v", err)
	}
	defer os.RemoveAll(tmpDir)

	gcsPath := "path/to/module/"
	server := fakestorage.NewServer(
		[]fakestorage.Object{
			{ObjectAttrs: fakestorage.ObjectAttrs{
				BucketName: "my-bucket", Name: "path/to/module/module.ko",
			},
				Content: []byte("content"),
			},
		})
	module := "module.ko"
	wantContent := "content"
	defer server.Stop()
	cfg := &GCSConfig{bucket: server.Client().Bucket("my-bucket"), path: gcsPath}
	modulePath, err := cfg.DownloadKernelModule(ctx, module, tmpDir)
	if err != nil {
		t.Errorf("TestDownloadKernelModuleSuccessful: failed to download module: %v", err)
	}
	wantPath := filepath.Join(tmpDir, module)
	content, err := os.ReadFile(wantPath)
	if err != nil {
		t.Errorf("TestDownloadKernelModuleSuccessful: failed to read file: %v", err)
	}
	if string(content) != wantContent {
		t.Errorf("TestDownloadKernelModuleSuccessful: expected content: %s\t got content: %s", wantContent, content)
	}
	if modulePath != wantPath {
		t.Errorf("TestDownloadKernelModuleSuccessful: expected path: %s\t got path: %s", wantPath, modulePath)
	}
}

func TestDownloadKernelModuleFileAlreadyExists(t *testing.T) {
	ctx := context.Background()
	module := "module.ko"
	tmpDir, err := os.MkdirTemp("", "test-gcs")
	if err != nil {
		t.Fatalf("Failed to create temp dir: %v", err)
	}
	defer os.RemoveAll(tmpDir)
	dst := filepath.Join(tmpDir, module)
	f, err := os.Create(dst)
	f.Close()
	if err != nil {
		t.Fatalf("Failed to create temp file: %v", err)
	}
	gcsPath := "path/to/module/"
	server := fakestorage.NewServer(
		[]fakestorage.Object{
			{ObjectAttrs: fakestorage.ObjectAttrs{
				BucketName: "my-bucket", Name: "path/to/module/module.ko",
			},
				Content: []byte("content"),
			},
		})
	defer server.Stop()
	cfg := &GCSConfig{bucket: server.Client().Bucket("my-bucket"), path: gcsPath}
	_, err = cfg.DownloadKernelModule(ctx, module, tmpDir)
	wantErrMessage := fmt.Sprintf("failed to download module (%s) to path (%s). File already exists", module, dst)
	if err == nil || err.Error() != wantErrMessage {
		t.Errorf("TestDownloadKernelModuleFileAlreadyExists: expected error: %s\t got error: %v", wantErrMessage, err.Error())
	}
}

func TestDownloadKernelModuleDoesNotExist(t *testing.T) {
	ctx := context.Background()
	module := "module.ko"
	tmpDir, err := os.MkdirTemp("", "test-gcs")
	if err != nil {
		t.Fatalf("Failed to create temp dir: %v", err)
	}
	defer os.RemoveAll(tmpDir)
	gcsPath := "path/to/module/"
	server := fakestorage.NewServer(
		[]fakestorage.Object{
			{ObjectAttrs: fakestorage.ObjectAttrs{
				BucketName: "my-bucket", Name: "path/to/module/unknown.ko",
			},
				Content: []byte("content"),
			},
		})
	defer server.Stop()
	cfg := &GCSConfig{bucket: server.Client().Bucket("my-bucket"), path: gcsPath}
	_, err = cfg.DownloadKernelModule(ctx, module, tmpDir)
	if err == nil || !strings.HasPrefix(err.Error(), "failed to read from module") {
		t.Errorf("TestDownloadKernelModuleDoesNotExist: expected error: failed to read from module \t got error: %v", err)
	}
}