3rd commit

internal/gitaly/storage/raftmgr/manager_test.go

@@ -18,6 +18,8 @@ import (
 	"gitlab.com/gitlab-org/gitaly/v16/internal/gitaly/storage/keyvalue/databasemgr"
 	"gitlab.com/gitlab-org/gitaly/v16/internal/gitaly/storage/storagemgr/partition/log"
 	"gitlab.com/gitlab-org/gitaly/v16/internal/gitaly/storage/wal"
+	"gitlab.com/gitlab-org/gitaly/v16/internal/helper"
+	logrus "gitlab.com/gitlab-org/gitaly/v16/internal/log"
 	"gitlab.com/gitlab-org/gitaly/v16/internal/testhelper"
 	"gitlab.com/gitlab-org/gitaly/v16/internal/testhelper/testcfg"
 	"gitlab.com/gitlab-org/gitaly/v16/proto/go/gitalypb"
@@ -1195,3 +1197,99 @@ func TestManager_NotImplementedLogMethods(t *testing.T) {
 	err = mgr.DeleteLogEntry(1)
 	require.ErrorContains(t, err, "raft manager does not support DeleteLogEntry")
 }
+
+func TestManager_StorageConnection(t *testing.T) {
+	t.Parallel()
+
+	ctx := testhelper.Context(t)
+	cfg := testcfg.Build(t)
+	logger := testhelper.NewLogger(t)
+
+	cfg.Raft = config.Raft{
+		Enabled:         true,
+		ElectionTicks:   5,
+		RTTMilliseconds: 100,
+		SnapshotDir:     testhelper.TempDir(t),
+	}
+
+	dbPath := testhelper.TempDir(t)
+	dbMgr, err := databasemgr.NewDBManager(
+		ctx,
+		cfg.Storages,
+		func(log logrus.Logger, path string) (keyvalue.Store, error) {
+			return keyvalue.NewBadgerStore(log, filepath.Join(dbPath, path))
+		},
+		helper.NewNullTickerFactory(),
+		logger,
+	)
+	require.NoError(t, err)
+	t.Cleanup(dbMgr.Close)
+
+	storageName := cfg.Storages[0].Name
+	db, err := dbMgr.GetDB(storageName)
+	require.NoError(t, err)
+
+	partitionID := storage.PartitionID(1)
+	stagingDir := testhelper.TempDir(t)
+	stateDir := testhelper.TempDir(t)
+	posTracker := log.NewPositionTracker()
+
+	raftStorage, err := NewStorage(
+		cfg.Raft,
+		logger,
+		storageName,
+		partitionID,
+		db,
+		stagingDir,
+		stateDir,
+		nil,
+		posTracker,
+		NewMetrics(),
+	)
+	require.NoError(t, err)
+
+	raftNode, err := NewNode(cfg, logger, dbMgr, nil)
+	require.NoError(t, err)
+
+	baseFactory := DefaultFactory(cfg.Raft)
+
+	// Create wrapped factory that connects managers to storage
+	wrappedFactory := RaftFactoryWithCallback(raftNode, logger, baseFactory)
+
+	manager, err := wrappedFactory(storageName, partitionID, raftStorage, logger)
+	require.NoError(t, err)
+	t.Cleanup(func() { require.NoError(t, manager.Close()) })
+
+	t.Run("manager has storage wrapper set", func(t *testing.T) {
+		wrapper := manager.GetStorageWrapper()
+		require.NotNil(t, wrapper, "storage wrapper should be set")
+
+		expectedWrapper, err := raftNode.GetStorageWrapper(storageName)
+		require.NoError(t, err)
+
+		require.Same(t, expectedWrapper, wrapper, "manager should point to correct storage wrapper")
+	})
+
+	t.Run("manager is registered in wrapper", func(t *testing.T) {
+		wrapper, err := raftNode.GetStorageWrapper(storageName)
+		require.NoError(t, err)
+
+		registeredManager, err := wrapper.GetManagerByPartition(partitionID, storageName)
+		require.NoError(t, err)
+		require.NotNil(t, registeredManager, "manager should be found in registry")
+
+		require.Same(t, manager, registeredManager, "registry should contain created manager instance")
+	})
+
+	t.Run("manager connection is bidirectional", func(t *testing.T) {
+		wrapper := manager.GetStorageWrapper()
+		require.NotNil(t, wrapper)
+
+		registeredManager, err := wrapper.GetManagerByPartition(partitionID, storageName)
+		require.NoError(t, err)
+
+		// Verify bidirectional connection
+		require.Same(t, manager, registeredManager,
+			"manager->wrapper->manager should return same instance")
+	})
+}