2424import org .apache .iotdb .itbase .category .AIClusterIT ;
2525import org .apache .iotdb .itbase .env .BaseEnv ;
2626
27- import org .junit .After ;
27+ import org .junit .AfterClass ;
2828import org .junit .Assert ;
29- import org .junit .Before ;
29+ import org .junit .BeforeClass ;
3030import org .junit .Test ;
3131import org .junit .experimental .categories .Category ;
3232import org .junit .runner .RunWith ;
3333
3434import java .sql .Connection ;
3535import java .sql .ResultSet ;
36- import java .sql .ResultSetMetaData ;
3736import java .sql .SQLException ;
3837import java .sql .Statement ;
3938
4443@ Category ({AIClusterIT .class })
4544public class AINodeClusterConfigIT {
4645
47- @ Before
48- public void setUp () throws Exception {
49- // Init 1C1D1A cluster environment
46+ @ BeforeClass
47+ public static void setUp () throws Exception {
5048 EnvFactory .getEnv ().initClusterEnvironment (1 , 1 );
5149 }
5250
53- @ After
54- public void tearDown () throws Exception {
51+ @ AfterClass
52+ public static void tearDown () throws Exception {
5553 EnvFactory .getEnv ().cleanClusterEnvironment ();
5654 }
5755
5856 @ Test
59- public void aiNodeRegisterAndRemoveTestInTree () throws SQLException {
57+ public void aiNodeRegisterAndRemoveTest () throws SQLException {
58+ String show_sql = "SHOW AINODES" ;
59+ String title = "NodeID,Status,InternalAddress,InternalPort" ;
60+
61+ // Verify AINode exists via both dialects before removal
6062 try (Connection connection = EnvFactory .getEnv ().getConnection (BaseEnv .TREE_SQL_DIALECT );
6163 Statement statement = connection .createStatement ()) {
62- aiNodeRegisterAndRemoveTest (statement );
64+ verifyAINodeExists (statement , show_sql , title );
65+ }
66+ try (Connection connection = EnvFactory .getEnv ().getConnection (BaseEnv .TABLE_SQL_DIALECT );
67+ Statement statement = connection .createStatement ()) {
68+ verifyAINodeExists (statement , show_sql , title );
6369 }
64- }
6570
66- @ Test
67- public void aiNodeRegisterAndRemoveTestInTable () throws SQLException {
71+ // Remove AINode
72+ try (Connection connection = EnvFactory .getEnv ().getConnection (BaseEnv .TREE_SQL_DIALECT );
73+ Statement statement = connection .createStatement ()) {
74+ statement .execute ("REMOVE AINODE" );
75+ waitForAINodeRemoval (statement , show_sql , title );
76+ }
77+
78+ // Verify removal is visible via table dialect as well
6879 try (Connection connection = EnvFactory .getEnv ().getConnection (BaseEnv .TABLE_SQL_DIALECT );
6980 Statement statement = connection .createStatement ()) {
70- aiNodeRegisterAndRemoveTest (statement );
81+ try (ResultSet resultSet = statement .executeQuery (show_sql )) {
82+ checkHeader (resultSet .getMetaData (), title );
83+ int count = 0 ;
84+ while (resultSet .next ()) {
85+ count ++;
86+ }
87+ assertEquals (0 , count );
88+ }
7189 }
7290 }
7391
74- private void aiNodeRegisterAndRemoveTest (Statement statement ) throws SQLException {
75- String show_sql = "SHOW AINODES" ;
76- String title = "NodeID,Status,InternalAddress,InternalPort" ;
77- try (ResultSet resultSet = statement .executeQuery (show_sql )) {
78- ResultSetMetaData resultSetMetaData = resultSet .getMetaData ();
79- checkHeader (resultSetMetaData , title );
92+ private static void verifyAINodeExists (Statement statement , String showSql , String title )
93+ throws SQLException {
94+ try (ResultSet resultSet = statement .executeQuery (showSql )) {
95+ checkHeader (resultSet .getMetaData (), title );
8096 int count = 0 ;
8197 while (resultSet .next ()) {
8298 assertEquals ("2" , resultSet .getString (1 ));
@@ -85,22 +101,23 @@ private void aiNodeRegisterAndRemoveTest(Statement statement) throws SQLExceptio
85101 }
86102 assertEquals (1 , count );
87103 }
88- String remove_sql = "REMOVE AINODE" ;
89- statement .execute (remove_sql );
104+ }
105+
106+ private static void waitForAINodeRemoval (Statement statement , String showSql , String title )
107+ throws SQLException {
90108 for (int retry = 0 ; retry < 500 ; retry ++) {
91- try (ResultSet resultSet = statement .executeQuery (show_sql )) {
92- ResultSetMetaData resultSetMetaData = resultSet .getMetaData ();
93- checkHeader (resultSetMetaData , title );
109+ try (ResultSet resultSet = statement .executeQuery (showSql )) {
110+ checkHeader (resultSet .getMetaData (), title );
94111 int count = 0 ;
95112 while (resultSet .next ()) {
96113 count ++;
97114 }
98115 if (count == 0 ) {
99- return ; // Successfully removed the AI node
116+ return ;
100117 }
101118 }
102119 try {
103- Thread .sleep (1000 ); // Wait before retrying
120+ Thread .sleep (1000 );
104121 } catch (InterruptedException e ) {
105122 Thread .currentThread ().interrupt ();
106123 }
0 commit comments