wandb

1
package observability_test
2

3
import (
4
	"fmt"
5
	"io/fs"
6
	"os"
7
	"testing"
8

9
	"github.com/spf13/afero"
10
	"github.com/stretchr/testify/assert"
11

12
	"github.com/wandb/wandb/core/pkg/observability"
13
)
14

15
type AferoFs struct {
16
	Afs afero.Fs
17
}
18

19
func (afs AferoFs) MkdirAll(path string, perm os.FileMode) error {
20
	return afs.Afs.MkdirAll(path, perm)
21
}
22

23
func (afs AferoFs) OpenFile(name string, flag int, perm os.FileMode) (fs.File, error) {
24
	return afs.Afs.OpenFile(name, flag, perm)
25
}
26

27
func TestGetLoggerPath(t *testing.T) {
28
	fs := AferoFs{Afs: afero.NewMemMapFs()}
29

30
	os.Setenv("WANDB_CACHE_DIR", "/tmp/wandb")
31
	defer os.Unsetenv("WANDB_CACHE_DIR")
32

33
	file, err := observability.GetLoggerPathFS(fs)
34
	assert.NoError(t, err)
35
	assert.NotNil(t, file)
36

37
	// Type assert to afero.File to access the Name method
38
	aferoFile, ok := file.(afero.File)
39
	assert.True(t, ok, "File should be of type afero.File")
40

41
	// Assert file was created
42
	_, err = fs.Afs.Stat(aferoFile.Name())
43
	assert.NoError(t, err)
44

45
	aferoFile.Close()
46
}
47

48
func TestGetLoggerPath_MkdirFail(t *testing.T) {
49
	fs := AferoFs{Afs: afero.NewReadOnlyFs(afero.NewMemMapFs())}
50

51
	os.Setenv("WANDB_CACHE_DIR", "/tmp/wandb")
52
	defer os.Unsetenv("WANDB_CACHE_DIR")
53

54
	_, err := observability.GetLoggerPathFS(fs)
55
	assert.Error(t, err, "Expected an error when failing to create a directory")
56
}
57

58
type FailOpenFileFs struct {
59
	afero.Fs
60
}
61

62
func (fs FailOpenFileFs) OpenFile(name string, flag int, perm os.FileMode) (afero.File, error) {
63
	return nil, fmt.Errorf("simulated open file error")
64
}
65

66
func TestGetLoggerPath_OpenFileFail(t *testing.T) {
67
	fs := AferoFs{Afs: FailOpenFileFs{afero.NewMemMapFs()}}
68

69
	os.Setenv("WANDB_CACHE_DIR", "/tmp/wandb")
70
	defer os.Unsetenv("WANDB_CACHE_DIR")
71

72
	_, err := observability.GetLoggerPathFS(fs)
73
	assert.Error(t, err, "Expected an error when failing to open a file")
74
}
75