@@ -206,9 +206,9 @@ class TPQTabletFixture : public NUnitTest::TBaseFixture {
206206 void StartPQWriteTxsObserver ();
207207 void WaitForPQWriteTxs ();
208208
209- template <class T > void WaitForEvent ();
210- void WaitForCalcPredicateResult ();
211- void WaitForProposePartitionConfigResult ();
209+ template <class T > void WaitForEvent (size_t count );
210+ void WaitForCalcPredicateResult (size_t count = 1 );
211+ void WaitForProposePartitionConfigResult (size_t count = 1 );
212212
213213 void TestWaitingForTEvReadSet (size_t senders, size_t receivers);
214214
@@ -524,14 +524,16 @@ void TPQTabletFixture::WaitDropTabletReply(const TDropTabletReplyMatcher& matche
524524}
525525
526526template <class T >
527- void TPQTabletFixture::WaitForEvent ()
527+ void TPQTabletFixture::WaitForEvent (size_t count )
528528{
529529 bool found = false ;
530+ size_t received = 0 ;
530531
531532 TTestActorRuntimeBase::TEventObserver prev;
532- auto observer = [&found, &prev](TAutoPtr<IEventHandle>& event) {
533+ auto observer = [&found, &prev, &received, count ](TAutoPtr<IEventHandle>& event) {
533534 if (auto * msg = event->CastAsLocal <T>()) {
534- found = true ;
535+ ++received;
536+ found = (received >= count);
535537 }
536538
537539 return prev ? prev (event) : TTestActorRuntimeBase::EEventAction::PROCESS;
@@ -549,14 +551,14 @@ void TPQTabletFixture::WaitForEvent()
549551 Ctx->Runtime ->SetObserverFunc (prev);
550552}
551553
552- void TPQTabletFixture::WaitForCalcPredicateResult ()
554+ void TPQTabletFixture::WaitForCalcPredicateResult (size_t count )
553555{
554- WaitForEvent<TEvPQ::TEvTxCalcPredicateResult>();
556+ WaitForEvent<TEvPQ::TEvTxCalcPredicateResult>(count );
555557}
556558
557- void TPQTabletFixture::WaitForProposePartitionConfigResult ()
559+ void TPQTabletFixture::WaitForProposePartitionConfigResult (size_t count )
558560{
559- WaitForEvent<TEvPQ::TEvProposePartitionConfigResult>();
561+ WaitForEvent<TEvPQ::TEvProposePartitionConfigResult>(count );
560562}
561563
562564std::unique_ptr<TEvPersQueue::TEvRequest> TPQTabletFixture::MakeGetOwnershipRequest (const TGetOwnershipRequestParams& params,
@@ -1370,7 +1372,7 @@ Y_UNIT_TEST_F(Read_TEvTxCommit_After_Restart, TPQTabletFixture)
13701372
13711373 WaitForCalcPredicateResult ();
13721374
1373- // the transaction is now in the WAIT_RS state on disk and in memory
1375+ // the transaction is now in the WAIT_RS state in memory and PLANNED state in disk
13741376
13751377 PQTabletRestart (*Ctx);
13761378
@@ -1416,7 +1418,145 @@ Y_UNIT_TEST_F(Config_TEvTxCommit_After_Restart, TPQTabletFixture)
14161418
14171419 WaitForProposePartitionConfigResult ();
14181420
1419- // the transaction is now in the WAIT_RS state on disk and in memory
1421+ // the transaction is now in the WAIT_RS state in memory and PLANNED state in disk
1422+
1423+ PQTabletRestart (*Ctx);
1424+
1425+ tablet->SendReadSet (*Ctx->Runtime , {.Step =100 , .TxId =txId, .Target =Ctx->TabletId , .Decision =NKikimrTx::TReadSetData::DECISION_COMMIT});
1426+
1427+ WaitProposeTransactionResponse ({.TxId =txId,
1428+ .Status =NKikimrPQ::TEvProposeTransactionResult::COMPLETE});
1429+
1430+ tablet->SendReadSetAck (*Ctx->Runtime , {.Step =100 , .TxId =txId, .Source =Ctx->TabletId });
1431+ WaitReadSetAck (*tablet, {.Step =100 , .TxId =txId, .Source =mockTabletId, .Target =Ctx->TabletId , .Consumer =Ctx->TabletId });
1432+ }
1433+
1434+ Y_UNIT_TEST_F (One_Tablet_For_All_Partitions, TPQTabletFixture)
1435+ {
1436+ const ui64 txId = 67890 ;
1437+
1438+ PQTabletPrepare ({.partitions =1 }, {}, *Ctx);
1439+
1440+ auto tabletConfig = NHelpers::MakeConfig ({.Version =2 ,
1441+ .Consumers ={
1442+ {.Consumer =" client-1" Generation =0 },
1443+ {.Consumer =" client-3" Generation =7 }
1444+ },
1445+ .Partitions ={
1446+ {.Id =0 },
1447+ {.Id =1 },
1448+ {.Id =2 }
1449+ },
1450+ .AllPartitions ={
1451+ {.Id =0 , .TabletId =Ctx->TabletId , .Children ={1 , 2 }, .Parents ={}},
1452+ {.Id =1 , .TabletId =Ctx->TabletId , .Children ={}, .Parents ={0 }},
1453+ {.Id =2 , .TabletId =Ctx->TabletId , .Children ={}, .Parents ={0 }}
1454+ }});
1455+
1456+ SendProposeTransactionRequest ({.TxId =txId,
1457+ .Configs =NHelpers::TConfigParams{
1458+ .Tablet =tabletConfig,
1459+ .Bootstrap =NHelpers::MakeBootstrapConfig (),
1460+ }});
1461+ WaitProposeTransactionResponse ({.TxId =txId,
1462+ .Status =NKikimrPQ::TEvProposeTransactionResult::PREPARED});
1463+
1464+ SendPlanStep ({.Step =100 , .TxIds ={txId}});
1465+
1466+ WaitForProposePartitionConfigResult (2 );
1467+
1468+ // the transaction is now in the WAIT_RS state in memory and PLANNED state in disk
1469+
1470+ PQTabletRestart (*Ctx);
1471+
1472+ WaitProposeTransactionResponse ({.TxId =txId,
1473+ .Status =NKikimrPQ::TEvProposeTransactionResult::COMPLETE});
1474+ }
1475+
1476+ Y_UNIT_TEST_F (One_New_Partition_In_Another_Tablet, TPQTabletFixture)
1477+ {
1478+ const ui64 txId = 67890 ;
1479+ const ui64 mockTabletId = 22222 ;
1480+
1481+ NHelpers::TPQTabletMock* tablet = CreatePQTabletMock (mockTabletId);
1482+ PQTabletPrepare ({.partitions =1 }, {}, *Ctx);
1483+
1484+ auto tabletConfig = NHelpers::MakeConfig ({.Version =2 ,
1485+ .Consumers ={
1486+ {.Consumer =" client-1" Generation =0 },
1487+ {.Consumer =" client-3" Generation =7 }
1488+ },
1489+ .Partitions ={
1490+ {.Id =0 },
1491+ {.Id =1 },
1492+ },
1493+ .AllPartitions ={
1494+ {.Id =0 , .TabletId =Ctx->TabletId , .Children ={1 , 2 }, .Parents ={}},
1495+ {.Id =1 , .TabletId =Ctx->TabletId , .Children ={}, .Parents ={0 }},
1496+ {.Id =2 , .TabletId =mockTabletId, .Children ={}, .Parents ={0 }}
1497+ }});
1498+
1499+ SendProposeTransactionRequest ({.TxId =txId,
1500+ .Configs =NHelpers::TConfigParams{
1501+ .Tablet =tabletConfig,
1502+ .Bootstrap =NHelpers::MakeBootstrapConfig (),
1503+ }});
1504+ WaitProposeTransactionResponse ({.TxId =txId,
1505+ .Status =NKikimrPQ::TEvProposeTransactionResult::PREPARED});
1506+
1507+ SendPlanStep ({.Step =100 , .TxIds ={txId}});
1508+
1509+ WaitForProposePartitionConfigResult (2 );
1510+
1511+ // the transaction is now in the WAIT_RS state in memory and PLANNED state in disk
1512+
1513+ PQTabletRestart (*Ctx);
1514+
1515+ tablet->SendReadSet (*Ctx->Runtime , {.Step =100 , .TxId =txId, .Target =Ctx->TabletId , .Decision =NKikimrTx::TReadSetData::DECISION_COMMIT});
1516+
1517+ WaitProposeTransactionResponse ({.TxId =txId,
1518+ .Status =NKikimrPQ::TEvProposeTransactionResult::COMPLETE});
1519+
1520+ tablet->SendReadSetAck (*Ctx->Runtime , {.Step =100 , .TxId =txId, .Source =Ctx->TabletId });
1521+ WaitReadSetAck (*tablet, {.Step =100 , .TxId =txId, .Source =mockTabletId, .Target =Ctx->TabletId , .Consumer =Ctx->TabletId });
1522+ }
1523+
1524+ Y_UNIT_TEST_F (All_New_Partitions_In_Another_Tablet, TPQTabletFixture)
1525+ {
1526+ const ui64 txId = 67890 ;
1527+ const ui64 mockTabletId = 22222 ;
1528+
1529+ NHelpers::TPQTabletMock* tablet = CreatePQTabletMock (mockTabletId);
1530+ PQTabletPrepare ({.partitions =1 }, {}, *Ctx);
1531+
1532+ auto tabletConfig = NHelpers::MakeConfig ({.Version =2 ,
1533+ .Consumers ={
1534+ {.Consumer =" client-1" Generation =0 },
1535+ {.Consumer =" client-3" Generation =7 }
1536+ },
1537+ .Partitions ={
1538+ {.Id =0 },
1539+ {.Id =1 },
1540+ },
1541+ .AllPartitions ={
1542+ {.Id =0 , .TabletId =Ctx->TabletId , .Children ={}, .Parents ={2 }},
1543+ {.Id =1 , .TabletId =Ctx->TabletId , .Children ={}, .Parents ={2 }},
1544+ {.Id =2 , .TabletId =mockTabletId, .Children ={0 , 1 }, .Parents ={}}
1545+ }});
1546+
1547+ SendProposeTransactionRequest ({.TxId =txId,
1548+ .Configs =NHelpers::TConfigParams{
1549+ .Tablet =tabletConfig,
1550+ .Bootstrap =NHelpers::MakeBootstrapConfig (),
1551+ }});
1552+ WaitProposeTransactionResponse ({.TxId =txId,
1553+ .Status =NKikimrPQ::TEvProposeTransactionResult::PREPARED});
1554+
1555+ SendPlanStep ({.Step =100 , .TxIds ={txId}});
1556+
1557+ WaitForProposePartitionConfigResult (2 );
1558+
1559+ // the transaction is now in the WAIT_RS state in memory and PLANNED state in disk
14201560
14211561 PQTabletRestart (*Ctx);
14221562
0 commit comments