fix: restore timeout in OmniSuite.SetupTest

The change from `context.WithTimeout` was unintentional, so lets revert that.

Signed-off-by: Dmitriy Matrenichev <[email protected]>

Author: Dmitriy Matrenichev

internal/backend/runtime/omni/controllers/omni/config_patch_cleanup_test.go

@@ -6,6 +6,7 @@
 package omni_test
 
 import (
+	"context"
 	"testing"
 	"testing/synctest"
 	"time"
@@ -27,6 +28,7 @@ type ConfigPatchCleanupSuite struct {
 }
 
 func (suite *ConfigPatchCleanupSuite) SetupTest() {
+	suite.OmniSuite.ctx, suite.OmniSuite.ctxCancel = context.WithCancel(context.Background()) //nolint:fatcontext
 	suite.OmniSuite.disableConnections = true
 	suite.OmniSuite.SetupTest()
 	suite.startRuntime()

internal/backend/runtime/omni/controllers/omni/etcd_backup.go

@@ -17,6 +17,7 @@ import (
 	"github.com/cosi-project/runtime/pkg/resource"
 	"github.com/cosi-project/runtime/pkg/safe"
 	"github.com/cosi-project/runtime/pkg/state"
+	"github.com/dustin/go-humanize"
 	"github.com/siderolabs/gen/channel"
 	"github.com/siderolabs/gen/containers"
 	machineapi "github.com/siderolabs/talos/pkg/machinery/api/machine"
@@ -387,6 +388,8 @@ func (ctrl *EtcdBackupController) doBackup(
 		return fmt.Errorf("failed to get store: %w", err)
 	}
 
+	crdr := &countingReader{rdr: rdr}
+
 	if err := st.Upload(
 		ctx,
 		etcdbackup.Description{
@@ -399,7 +402,7 @@ func (ctrl *EtcdBackupController) doBackup(
 				EncryptionKey:             backupData.TypedSpec().Value.EncryptionKey,
 			},
 		},
-		rdr,
+		crdr,
 	); err != nil {
 		return fmt.Errorf("failed to upload etcd snapshot for cluster: %w", err)
 	}
@@ -410,6 +413,7 @@ func (ctrl *EtcdBackupController) doBackup(
 		"uploaded etcd snapshot",
 		zap.String("cluster_uuid", backupData.TypedSpec().Value.ClusterUuid),
 		zap.String("snapshot_name", etcdbackup.CreateSnapshotName(now)),
+		zap.String("snapshot_size", humanize.IBytes(uint64(crdr.total))),
 		zap.Time("ts", now),
 	)
 
@@ -460,3 +464,17 @@ func (ctrl *EtcdBackupController) updateBackupStatus(
 type TalosClient interface {
 	EtcdSnapshot(ctx context.Context, req *machineapi.EtcdSnapshotRequest, callOptions ...grpc.CallOption) (io.ReadCloser, error)
 }
+
+type countingReader struct {
+	rdr   io.Reader
+	total int
+}
+
+func (r *countingReader) Read(p []byte) (int, error) {
+	n, err := r.rdr.Read(p)
+	if n > 0 {
+		r.total += n
+	}
+
+	return n, err
+}

internal/backend/runtime/omni/controllers/omni/etcd_backup_test.go

@@ -57,9 +57,7 @@ import (
 func TestEtcdBackupControllerSuite(t *testing.T) {
 	t.Parallel()
 
-	synctest.Run(func() {
-		suite.Run(t, new(EtcdBackupControllerSuite))
-	})
+	synctest.Run(func() { suite.Run(t, new(EtcdBackupControllerSuite)) })
 }
 
 const description = "test"
@@ -82,6 +80,8 @@ func (suite *EtcdBackupControllerSuite) register2(ctrl controller.Controller, er
 }
 
 func (suite *EtcdBackupControllerSuite) SetupTest() {
+	suite.OmniSuite.ctx, suite.OmniSuite.ctxCancel = context.WithCancel(context.Background()) //nolint:fatcontext
+
 	suite.OmniSuite.disableConnections = true
 
 	suite.OmniSuite.SetupTest()
@@ -118,7 +118,7 @@ func (suite *EtcdBackupControllerSuite) TestEtcdBackup() {
 	clustersData := createClusters(suite, clusterNames, time.Hour)
 	start := time.Now()
 
-	time.Sleep(11 * time.Minute)
+	// Wait for the first backup to be created.
 	synctest.Wait()
 
 	rtestutils.AssertResources(

internal/backend/runtime/omni/controllers/omni/omni_test.go

@@ -406,7 +406,15 @@ func (suite *OmniSuite) newServerWithTalosVersion(suffix, talosVersion string) (
 }
 
 func (suite *OmniSuite) SetupTest() {
-	suite.ctx, suite.ctxCancel = context.WithCancel(suite.T().Context())
+	if suite.ctx == nil {
+		// This check is important, because some of our tests use synctest package, which
+		// "mocks" the time, and advances it in multiple hours. Because the context is
+		// created inside the "bubble" when 20 seconds elapse, the suite will shut down
+		// before the test is actually finished. There are two test suites that currently
+		// need their separate contexts: [ConfigPatchCleanupSuite] and
+		// [EtcdBackupControllerSuite] in which we create context explicitly.
+		suite.ctx, suite.ctxCancel = context.WithTimeout(suite.T().Context(), 20*time.Second)
+	}
 
 	suite.stateBuilder = dynamicStateBuilder{m: map[resource.Namespace]state.CoreState{}}
 
@@ -484,6 +492,8 @@ func (suite *OmniSuite) TearDownTest() {
 	for _, s := range suite.grpcServers {
 		s.Stop()
 	}
+
+	suite.ctx, suite.ctxCancel = nil, nil
 }
 
 func assertResource[R rtestutils.ResourceWithRD](suite *OmniSuite, md interface{ ID() resource.ID }, assertionFunc func(r R, assertion *assert.Assertions)) {