cost-model/core/pkg/storage/test.go

package storage

import (
	"bytes"
	"fmt"
	"io"
	"os"
	"path"
	"path/filepath"
	"strings"
	"testing"

	"github.com/opencost/opencost/core/pkg/util/json"
)

type testFileContent struct {
	Field1 int    `json:"field_1"`
	Field2 string `json:"field_2"`
}

var tfc = testFileContent{
	Field1: 101,
	Field2: "TEST_FILE_CONTENT",
}

const testpath = "opencost/storage/"

func createFiles(files []string, testName string, store Storage) error {
	b, err := json.Marshal(tfc)
	if err != nil {
		return fmt.Errorf("failed to marshal file content: %w", err)
	}

	for _, fileName := range files {
		filePath := path.Join(testpath, testName, fileName)
		err = store.Write(filePath, b)
		if err != nil {
			return fmt.Errorf("failed to write file '%s': %w ", filePath, err)
		}
	}

	return nil
}

func cleanupFiles(files []string, testName string, store Storage) error {
	for _, fileName := range files {
		filePath := path.Join(testpath, testName, fileName)
		err := store.Remove(filePath)
		if err != nil {
			return fmt.Errorf("failed to remove file '%s': %w ", filePath, err)
		}
	}

	return nil
}

func TestStorageList(t *testing.T, store Storage) {
	testName := "list"

	fileNames := []string{
		"/file0.json",
		"/file1.json",
		"/dir0/file2.json",
		"/dir0/file3.json",
	}

	err := createFiles(fileNames, testName, store)
	if err != nil {
		t.Errorf("failed to create files: %s", err)
	}

	defer func() {
		err = cleanupFiles(fileNames, testName, store)
		if err != nil {
			t.Errorf("failed to clean up files: %s", err)
		}
	}()

	testCases := map[string]struct {
		path      string
		expected  []string
		expectErr bool
	}{
		"base dir files": {
			path: path.Join(testpath, testName),
			expected: []string{
				"file0.json",
				"file1.json",
			},
			expectErr: false,
		},
		"single nested dir files": {
			path: path.Join(testpath, testName, "dir0"),
			expected: []string{
				"file2.json",
				"file3.json",
			},
			expectErr: false,
		},
		"nonexistent dir files": {
			path:      path.Join(testpath, testName, "dir1"),
			expected:  []string{},
			expectErr: false,
		},
	}

	for name, tc := range testCases {
		t.Run(name, func(t *testing.T) {
			fileList, err := store.List(tc.path)
			if tc.expectErr == (err == nil) {
				if tc.expectErr {
					t.Errorf("expected error was not thrown")
					return
				}
				t.Errorf("unexpected error: %s", err.Error())
				return
			}

			if len(fileList) != len(tc.expected) {
				t.Errorf("file list length does not match expected length, actual: %d, expected: %d", len(fileList), len(tc.expected))
			}

			expectedSet := map[string]struct{}{}
			for _, expName := range tc.expected {
				expectedSet[expName] = struct{}{}
			}

			for _, file := range fileList {
				_, ok := expectedSet[file.Name]
				if !ok {
					t.Errorf("unexpect file in list %s", file.Name)
				}

				if file.Size == 0 {
					t.Errorf("file size is not set")
				}

				if file.ModTime.IsZero() {
					t.Errorf("file mod time is not set")
				}
			}
		})
	}
}

func TestStorageListDirectories(t *testing.T, store Storage) {
	testName := "list_directories"

	fileNames := []string{
		"/file0.json",
		"/dir0/file2.json",
		"/dir0/file3.json",
		"/dir0/dir1/file4.json",
		"/dir0/dir2/file5.json",
	}

	err := createFiles(fileNames, testName, store)
	if err != nil {
		t.Errorf("failed to create files: %s", err)
	}

	defer func() {
		err = cleanupFiles(fileNames, testName, store)
		if err != nil {
			t.Errorf("failed to clean up files: %s", err)
		}
	}()

	testCases := map[string]struct {
		path      string
		expected  []string
		expectErr bool
	}{
		"root dir dir": {
			path: "",
			expected: []string{
				strings.Split(testpath, "/")[0] + "/",
			},
			expectErr: false,
		},
		"base dir dir": {
			path: path.Join(testpath, testName),
			expected: []string{
				path.Join(testpath, testName, "dir0") + "/",
			},
			expectErr: false,
		},
		"single nested dir files": {
			path: path.Join(testpath, testName, "dir0"),
			expected: []string{
				path.Join(testpath, testName, "dir0", "dir1") + "/",
				path.Join(testpath, testName, "dir0", "dir2") + "/",
			},
			expectErr: false,
		},
		"dir with no sub dirs": {
			path:      path.Join(testpath, testName, "dir0/dir1"),
			expected:  []string{},
			expectErr: false,
		},
		"non-existent dir": {
			path:      path.Join(testpath, testName, "dir1"),
			expected:  []string{},
			expectErr: false,
		},
	}

	for name, tc := range testCases {
		t.Run(name, func(t *testing.T) {
			dirList, err := store.ListDirectories(tc.path)
			if tc.expectErr == (err == nil) {
				if tc.expectErr {
					t.Errorf("expected error was not thrown")
					return
				}
				t.Errorf("unexpected error: %s", err.Error())
				return
			}

			if len(dirList) != len(tc.expected) {
				t.Errorf("dir list length does not match expected length, actual: %d, expected: %d", len(dirList), len(tc.expected))
			}

			expectedSet := map[string]struct{}{}
			for _, expName := range tc.expected {
				expectedSet[expName] = struct{}{}
			}

			for _, dir := range dirList {
				_, ok := expectedSet[dir.Name]
				if !ok {
					t.Errorf("unexpect dir in list %s", dir.Name)
				}
			}
		})
	}
}

func TestStorageExists(t *testing.T, store Storage) {
	testName := "exists"

	fileNames := []string{
		"/file0.json",
	}

	err := createFiles(fileNames, testName, store)
	if err != nil {
		t.Errorf("failed to create files: %s", err)
	}

	defer func() {
		err = cleanupFiles(fileNames, testName, store)
		if err != nil {
			t.Errorf("failed to clean up files: %s", err)
		}
	}()

	testCases := map[string]struct {
		path      string
		expected  bool
		expectErr bool
	}{
		"file exists": {
			path:      path.Join(testpath, testName, "file0.json"),
			expected:  true,
			expectErr: false,
		},
		"file does not exist": {
			path:      path.Join(testpath, testName, "file1.json"),
			expected:  false,
			expectErr: false,
		},
		"dir does not exist": {
			path:      path.Join(testpath, testName, "dir0/file.json"),
			expected:  false,
			expectErr: false,
		},
	}

	for name, tc := range testCases {
		t.Run(name, func(t *testing.T) {
			exists, err := store.Exists(tc.path)
			if tc.expectErr == (err == nil) {
				if tc.expectErr {
					t.Errorf("expected error was not thrown")
					return
				}
				t.Errorf("unexpected error: %s", err.Error())
				return
			}

			if exists != tc.expected {
				t.Errorf("file exists output did not match expected")
			}
		})
	}
}

func TestStorageRead(t *testing.T, store Storage) {
	testName := "read"

	fileNames := []string{
		"/file0.json",
	}

	err := createFiles(fileNames, testName, store)
	if err != nil {
		t.Errorf("failed to create files: %s", err)
	}

	defer func() {
		err = cleanupFiles(fileNames, testName, store)
		if err != nil {
			t.Errorf("failed to clean up files: %s", err)
		}
	}()

	testCases := map[string]struct {
		path      string
		expectErr bool
	}{
		"file exists": {
			path:      path.Join(testpath, testName, "file0.json"),
			expectErr: false,
		},
		"file does not exist": {
			path:      path.Join(testpath, testName, "file1.json"),
			expectErr: true,
		},
		"dir does not exist": {
			path:      path.Join(testpath, testName, "dir0/file.json"),
			expectErr: true,
		},
	}

	for name, tc := range testCases {
		t.Run(name, func(t *testing.T) {
			b, err := store.Read(tc.path)
			if tc.expectErr && err != nil {
				return
			}
			if tc.expectErr == (err == nil) {
				if tc.expectErr {
					t.Errorf("expected error was not thrown")
					return
				}
				t.Errorf("unexpected error: %s", err.Error())
				return
			}
			var content testFileContent
			err = json.Unmarshal(b, &content)
			if err != nil {
				t.Errorf("could not unmarshal file content")
				return
			}

			if content != tfc {
				t.Errorf("file content did not match writen value")
			}
		})
	}
}

func TestStorageStat(t *testing.T, store Storage) {
	testName := "stat"

	fileNames := []string{
		"/file0.json",
	}

	err := createFiles(fileNames, testName, store)
	if err != nil {
		t.Errorf("failed to create files: %s", err)
	}

	defer func() {
		err = cleanupFiles(fileNames, testName, store)
		if err != nil {
			t.Errorf("failed to clean up files: %s", err)
		}
	}()

	testCases := map[string]struct {
		path      string
		expected  *StorageInfo
		expectErr bool
	}{
		"base dir": {
			path: path.Join(testpath, testName, "file0.json"),
			expected: &StorageInfo{
				Name: "file0.json",
				Size: 45,
			},
			expectErr: false,
		},
		"file does not exist": {
			path:      path.Join(testpath, testName, "file1.json"),
			expected:  nil,
			expectErr: true,
		},
	}

	for name, tc := range testCases {
		t.Run(name, func(t *testing.T) {
			status, err := store.Stat(tc.path)
			if tc.expectErr && err != nil {
				return
			}
			if tc.expectErr == (err == nil) {
				if tc.expectErr {
					t.Errorf("expected error was not thrown")
					return
				}
				t.Errorf("unexpected error: %s", err.Error())
				return
			}

			if status.Name != tc.expected.Name {
				t.Errorf("status name did name match expected, actual: %s, expected: %s", status.Name, tc.expected.Name)
			}

			if status.Size != tc.expected.Size {
				t.Errorf("status name did size match expected, actual: %d, expected: %d", status.Size, tc.expected.Size)
			}

			if status.ModTime.IsZero() {
				t.Errorf("status mod time is not set")
			}

		})
	}
}

func TestStorageReadToLocalFile(t *testing.T, store Storage) {
	testName := "read_to_local_file"

	fileNames := []string{
		"/file0.json",
	}

	err := createFiles(fileNames, testName, store)
	if err != nil {
		t.Fatalf("failed to create files: %s", err)
	}

	defer func() {
		err = cleanupFiles(fileNames, testName, store)
		if err != nil {
			t.Fatalf("failed to clean up files: %s", err)
		}
	}()

	testCases := map[string]struct {
		path      string
		expectErr bool
	}{
		"file exists": {
			path:      path.Join(testpath, testName, "file0.json"),
			expectErr: false,
		},
		"file does not exist": {
			path:      path.Join(testpath, testName, "file1.json"),
			expectErr: true,
		},
		"dir does not exist": {
			path:      path.Join(testpath, testName, "dir0/file.json"),
			expectErr: true,
		},
	}

	for name, tc := range testCases {
		t.Run(name, func(t *testing.T) {
			destPath := filepath.Join(t.TempDir(), "out.json")

			err := store.ReadToLocalFile(tc.path, destPath)
			if tc.expectErr {
				if err == nil {
					t.Fatalf("expected error was not thrown")
				}
				return
			}

			if err != nil {
				t.Fatalf("unexpected error: %s", err.Error())
			}

			b, err := os.ReadFile(destPath)
			if err != nil {
				t.Fatalf("reading destination file: %s", err)
			}

			var content testFileContent
			err = json.Unmarshal(b, &content)
			if err != nil {
				t.Fatalf("could not unmarshal file content: %s", err)
			}

			if content != tfc {
				t.Fatalf("file content did not match written value")
			}
		})
	}
}

func TestStorageWriteStream(t *testing.T, store Storage) {
	testName := "write_stream"

	testCases := map[string]struct {
		path     string
		chunks   [][]byte
		prewrite bool
	}{
		"single chunk": {
			path:   path.Join(testpath, testName, "single.bin"),
			chunks: [][]byte{[]byte("single chunk data")},
		},
		"multi-chunk": {
			path:   path.Join(testpath, testName, "multi.bin"),
			chunks: [][]byte{[]byte("alpha"), []byte("beta"), []byte("gamma")},
		},
		"nested path": {
			path:   path.Join(testpath, testName, "sub/dir/data.bin"),
			chunks: [][]byte{[]byte("nested data")},
		},
		"empty content": {
			path:   path.Join(testpath, testName, "empty.bin"),
			chunks: [][]byte{},
		},
		"overwrite existing": {
			path:     path.Join(testpath, testName, "overwrite.bin"),
			chunks:   [][]byte{[]byte("replaced content")},
			prewrite: true,
		},
	}

	for name, tc := range testCases {
		t.Run(name, func(t *testing.T) {
			defer store.Remove(tc.path)

			if tc.prewrite {
				b, err := json.Marshal(tfc)
				if err != nil {
					t.Fatalf("marshal fixture: %s", err)
				}
				if err = store.Write(tc.path, b); err != nil {
					t.Fatalf("pre-write: %s", err)
				}
			}

			var want []byte
			for _, chunk := range tc.chunks {
				want = append(want, chunk...)
			}

			w, err := store.WriteStream(tc.path)
			if err != nil {
				t.Fatalf("WriteStream: %s", err)
			}

			for _, chunk := range tc.chunks {
				n, writeErr := w.Write(chunk)
				if writeErr != nil {
					_ = w.Close()
					t.Fatalf("Write: %s", writeErr)
				}
				if n != len(chunk) {
					_ = w.Close()
					t.Fatalf("short write: wrote %d of %d bytes", n, len(chunk))
				}
			}

			if err = w.Close(); err != nil {
				t.Fatalf("Close: %s", err)
			}

			got, err := store.Read(tc.path)
			if err != nil {
				t.Fatalf("Read after WriteStream: %s", err)
			}

			if !bytes.Equal(got, want) {
				t.Errorf("content mismatch: got %q, want %q", got, want)
			}
		})
	}
}

func TestStorageReadStream(t *testing.T, store Storage) {
	testName := "read_stream"

	fileNames := []string{
		"/file0.json",
	}

	err := createFiles(fileNames, testName, store)
	if err != nil {
		t.Fatalf("failed to create files: %s", err)
	}

	defer func() {
		err = cleanupFiles(fileNames, testName, store)
		if err != nil {
			t.Fatalf("failed to clean up files: %s", err)
		}
	}()

	testCases := map[string]struct {
		path      string
		expectErr bool
	}{
		"file exists": {
			path:      path.Join(testpath, testName, "file0.json"),
			expectErr: false,
		},
		"file does not exist": {
			path:      path.Join(testpath, testName, "file1.json"),
			expectErr: true,
		},
		"dir does not exist": {
			path:      path.Join(testpath, testName, "dir0/file.json"),
			expectErr: true,
		},
	}

	for name, tc := range testCases {
		t.Run(name, func(t *testing.T) {
			r, err := store.ReadStream(tc.path)
			if tc.expectErr {
				if err == nil {
					if r != nil {
						_ = r.Close()
					}
					t.Fatalf("expected error was not thrown")
				}
				return
			}

			if err != nil {
				t.Fatalf("unexpected error: %s", err.Error())
			}
			defer r.Close()

			b, err := io.ReadAll(r)
			if err != nil {
				t.Fatalf("reading stream: %s", err)
			}

			var content testFileContent
			err = json.Unmarshal(b, &content)
			if err != nil {
				t.Fatalf("could not unmarshal file content: %s", err)
			}

			if content != tfc {
				t.Fatalf("file content did not match written value")
			}
		})
	}
}