Merge pull request #39 from Photoroom/ben/adding_fs_basic_unit_test

[CI] Basic filesystem path unit test

README.md

@@ -98,7 +98,7 @@ $ go run -race cmd/main/main.go
 From the root folder
 
 ```bash
-$ go test -v tests/client_test.go
+$ go test -v tests/*
 ```
 
 ## Refresh the python package and its binaries

tests/client_test.go → tests/client_db_test.go

@@ -289,4 +289,32 @@ func TestRanks(t *testing.T) {
 	client_1.Stop()
 }
 
+func TestTags(t *testing.T) {
+	clientConfig := get_default_test_config()
+	clientConfig.SamplesBufferSize = 1
+
+	dbConfig := clientConfig.SourceConfig.(datago.SourceDBConfig)
+	dbConfig.Tags = "v4_trainset_hq"
+	clientConfig.SourceConfig = dbConfig
+
+	client := datago.GetClient(clientConfig)
+
+	countains := func(slice []string, element string) bool {
+		for _, item := range slice {
+			if item == element {
+				return true
+			}
+		}
+		return false
+	}
+
+	for i := 0; i < 10; i++ {
+		sample := client.GetSample()
+		if !countains(sample.Tags, "v4_trainset_hq") {
+			t.Errorf("Sample is missing the required tag")
+		}
+	}
+	client.Stop()
+}
+
 // FIXME: Could do with a lot of tests on the filesystem side

tests/client_filesystem_test.go

@@ -0,0 +1,95 @@
+package datago_test
+
+import (
+	datago "datago/pkg"
+	"fmt"
+	"image"
+	"image/color"
+	"image/png"
+	"math/rand"
+	"os"
+	"testing"
+	"time"
+)
+
+// Function to generate and save a random image
+func generateRandomImage(width, height int, filename string) error {
+	// Create a new RGBA image
+	img := image.NewRGBA(image.Rect(0, 0, width, height))
+
+	// Seed the random number generator
+	rand.Seed(time.Now().UnixNano())
+
+	// Fill the image with random colors
+	for y := 0; y < height; y++ {
+		for x := 0; x < width; x++ {
+			r := uint8(rand.Intn(256))
+			g := uint8(rand.Intn(256))
+			b := uint8(rand.Intn(256))
+			img.Set(x, y, color.RGBA{R: r, G: g, B: b, A: 255})
+		}
+	}
+
+	// Create the output file
+	file, err := os.Create(filename)
+	if err != nil {
+		return err
+	}
+	defer file.Close()
+
+	// Encode the image to the file in PNG format
+	err = png.Encode(file, img)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func init_test_data(num_images int) string {
+	// Get a temporary directory
+	base_tmp_dir := os.TempDir()
+	dir := base_tmp_dir + "/datago_test"
+	err := os.MkdirAll(dir, os.ModePerm)
+	if err != nil {
+		panic(err)
+	}
+
+	// Generate random images and save them to the temporary directory
+	for i := 0; i < num_images; i++ {
+		filename := dir + "/test_image_" + fmt.Sprint(i) + ".png"
+		err := generateRandomImage(1024, 1024, filename)
+		if err != nil {
+			panic(err)
+		}
+	}
+
+	return dir
+}
+
+func TestFilesystemLoad(t *testing.T) {
+	// Initialize the test data
+	num_images := 10
+
+	test_dir := init_test_data(num_images)
+	defer os.RemoveAll(test_dir)
+
+	// Set the environment variable
+	os.Setenv("DATAGO_TEST_FILESYSTEM", test_dir)
+
+	// Run the tests
+	config := datago.GetDatagoConfig()
+	fs_config := datago.GetSourceFileSystemConfig()
+	fs_config.RootPath = test_dir
+	config.SourceConfig = fs_config
+
+	client := datago.GetClient(config)
+	client.Start()
+
+	for i := 0; i < num_images; i++ {
+		sample := client.GetSample()
+		if sample.ID == "" {
+			t.Errorf("GetSample returned an unexpected error")
+		}
+	}
+}