Add tests

builder/store/database/lfs_meta_object_test.go

@@ -84,3 +84,101 @@
 	require.NotNil(t, err)
 
 }
+
+func TestLfsMetaStore_UpdateXnetUsed(t *testing.T) {
+	db := tests.InitTestDB()
+	defer db.Close()
+	ctx := context.TODO()
+
+	store := database.NewLfsMetaObjectStoreWithDB(db)
+
+	// Create a test LFS object
+	_, err := store.Create(ctx, database.LfsMetaObject{
+		RepositoryID: 123,
+		Oid:          "test-oid-123",
+		Size:         1024,
+		XnetUsed:     false,
+	})
+	require.Nil(t, err)
+
+	// Verify initial state
+	obj, err := store.FindByOID(ctx, 123, "test-oid-123")
+	require.Nil(t, err)
+	require.Equal(t, "test-oid-123", obj.Oid)
+	require.Equal(t, false, obj.XnetUsed)
+
+	// Update XnetUsed to true
+	err = store.UpdateXnetUsed(ctx, 123, "test-oid-123", true)
+	require.Nil(t, err)
+
+	// Verify update
+	obj, err = store.FindByOID(ctx, 123, "test-oid-123")
+	require.Nil(t, err)
+	require.Equal(t, true, obj.XnetUsed)
+
+	// Update XnetUsed back to false
+	err = store.UpdateXnetUsed(ctx, 123, "test-oid-123", false)
+	require.Nil(t, err)
+
+	// Verify update
+	obj, err = store.FindByOID(ctx, 123, "test-oid-123")
+	require.Nil(t, err)
+	require.Equal(t, false, obj.XnetUsed)
+
+	// Test updating non-existent object (should not error but affect 0 rows)
+	err = store.UpdateXnetUsed(ctx, 999, "non-existent-oid", true)
+	require.Nil(t, err)
+}
+
+func TestLfsMetaStore_CheckIfAllMigratedToXnet(t *testing.T) {
+	db := tests.InitTestDB()
+	defer db.Close()
+	ctx := context.TODO()
+
+	store := database.NewLfsMetaObjectStoreWithDB(db)
+
+	// Test with no objects
+	allMigrated, err := store.CheckIfAllMigratedToXnet(ctx, 123)
+	require.Nil(t, err)
+	require.True(t, allMigrated, "No objects should return true")
+
+	// Create test objects
+	testCases := []database.LfsMetaObject{
+		{RepositoryID: 123, Oid: "oid-1", Size: 1024, XnetUsed: false},
+		{RepositoryID: 123, Oid: "oid-2", Size: 2048, XnetUsed: false},
+		{RepositoryID: 456, Oid: "oid-3", Size: 3072, XnetUsed: true}, // Different repo
+	}
+
+	for _, tc := range testCases {
+		_, err := store.Create(ctx, tc)
+		require.Nil(t, err)
+	}
+
+	// Test with some objects not migrated
+	allMigrated, err = store.CheckIfAllMigratedToXnet(ctx, 123)
+	require.Nil(t, err)
+	require.False(t, allMigrated, "Should return false when some objects are not migrated")
+
+	// Test with all objects migrated in another repo
+	allMigrated, err = store.CheckIfAllMigratedToXnet(ctx, 456)
+	require.Nil(t, err)
+	require.True(t, allMigrated, "Should return true when all objects are migrated")
+
+	// Update one object to migrated
+	err = store.UpdateXnetUsed(ctx, 123, "oid-1", true)
+	require.Nil(t, err)
+
+	// Test with one migrated, one not migrated
+	allMigrated, err = store.CheckIfAllMigratedToXnet(ctx, 123)
+	require.Nil(t, err)
+	require.False(t, allMigrated, "Should return false when not all objects are migrated")
+
+	// Update all objects to migrated
+	err = store.UpdateXnetUsed(ctx, 123, "oid-2", true)
+	require.Nil(t, err)
+
+	// Test with all objects migrated
+	allMigrated, err = store.CheckIfAllMigratedToXnet(ctx, 123)
+	require.Nil(t, err)
+	require.True(t, allMigrated, "Should return true when all objects are migrated")
+}

builder/store/database/xnet_migration_task_test.go

@@ -0,0 +1,82 @@
+package database_test
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+	"opencsg.com/csghub-server/builder/store/database"
+	"opencsg.com/csghub-server/common/tests"
+	"opencsg.com/csghub-server/common/types"
+)
+
+func TestXnetMigrationTaskStore_CRUD(t *testing.T) {
+	db := tests.InitTestDB()
+	defer db.Close()
+	ctx := context.TODO()
+
+	store := database.NewXnetMigrationTaskStoreWithDB(db)
+
+	// Test CreateXnetMigrationTask
+	err := store.CreateXnetMigrationTask(ctx, 123, "Initial task")
+	require.Nil(t, err)
+
+	// Test ListXnetMigrationTasksByRepoID
+	tasks, err := store.ListXnetMigrationTasksByRepoID(ctx, 123)
+	require.Nil(t, err)
+	require.Len(t, tasks, 1)
+	require.Equal(t, int64(123), tasks[0].RepositoryID)
+	require.Equal(t, "Initial task", tasks[0].LastMessage)
+	require.Equal(t, types.XnetMigrationTaskStatusPending, tasks[0].Status)
+
+	taskID := tasks[0].ID
+
+	// Test GetXnetMigrationTaskByID
+	task, err := store.GetXnetMigrationTaskByID(ctx, taskID)
+	require.Nil(t, err)
+	require.Equal(t, taskID, task.ID)
+	require.Equal(t, int64(123), task.RepositoryID)
+
+	// Test UpdateXnetMigrationTask
+	err = store.UpdateXnetMigrationTask(ctx, taskID, "Updated message", types.XnetMigrationTaskStatusRunning)
+	require.Nil(t, err)
+
+	// Verify update
+	task, err = store.GetXnetMigrationTaskByID(ctx, taskID)
+	require.Nil(t, err)
+	require.Equal(t, "Updated message", task.LastMessage)
+	require.Equal(t, types.XnetMigrationTaskStatusRunning, task.Status)
+
+	// Create more tasks with different statuses
+	err = store.CreateXnetMigrationTask(ctx, 456, "Task 2")
+	require.Nil(t, err)
+	tasks, err = store.ListXnetMigrationTasksByRepoID(ctx, 456)
+	require.Nil(t, err)
+	require.Len(t, tasks, 1)
+	err = store.UpdateXnetMigrationTask(ctx, tasks[0].ID, "Task 2 completed", types.XnetMigrationTaskStatusCompleted)
+	require.Nil(t, err)
+
+	// Test ListXnetMigrationTasksByStatus
+	runningTasks, err := store.ListXnetMigrationTasksByStatus(ctx, types.XnetMigrationTaskStatusRunning)
+	require.Nil(t, err)
+	require.Len(t, runningTasks, 1)
+	require.Equal(t, "Updated message", runningTasks[0].LastMessage)
+
+	completedTasks, err := store.ListXnetMigrationTasksByStatus(ctx, types.XnetMigrationTaskStatusCompleted)
+	require.Nil(t, err)
+	require.Len(t, completedTasks, 1)
+	require.Equal(t, "Task 2 completed", completedTasks[0].LastMessage)
+
+	pendingTasks, err := store.ListXnetMigrationTasksByStatus(ctx, types.XnetMigrationTaskStatusPending)
+	require.Nil(t, err)
+	require.Len(t, pendingTasks, 0)
+
+	// Test with non-existent repo ID
+	noTasks, err := store.ListXnetMigrationTasksByRepoID(ctx, 999)
+	require.Nil(t, err)
+	require.Len(t, noTasks, 0)
+
+	// Test with non-existent task ID
+	_, err = store.GetXnetMigrationTaskByID(ctx, 999)
+	require.NotNil(t, err)
+}

component/lfs.go

@@ -0,0 +1,13 @@
+package component
+
+import (
+	"context"
+
+	"opencsg.com/csghub-server/builder/store/database"
+)
+
+type LfsComponent interface {
+	DispatchLfsXnetProgress() error
+	DispatchLfsXnetResult() error
+	PublishLfsMigrationMessage(ctx context.Context, repo *database.Repository, oid string) error
+}