1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18 package org.apache.shardingsphere.infra.database.postgresql.checker;
19
20 import org.apache.shardingsphere.infra.database.core.checker.PrivilegeCheckType;
21 import org.apache.shardingsphere.infra.database.core.exception.MissingRequiredPrivilegeException;
22 import org.junit.jupiter.api.BeforeEach;
23 import org.junit.jupiter.api.Test;
24 import org.junit.jupiter.api.extension.ExtendWith;
25 import org.mockito.Mock;
26 import org.mockito.junit.jupiter.MockitoExtension;
27
28 import javax.sql.DataSource;
29 import java.sql.Connection;
30 import java.sql.DatabaseMetaData;
31 import java.sql.PreparedStatement;
32 import java.sql.ResultSet;
33 import java.sql.SQLException;
34
35 import static org.junit.jupiter.api.Assertions.assertThrows;
36 import static org.mockito.ArgumentMatchers.anyString;
37 import static org.mockito.Mockito.atLeastOnce;
38 import static org.mockito.Mockito.verify;
39 import static org.mockito.Mockito.when;
40
41 @ExtendWith(MockitoExtension.class)
42 class PostgreSQLDatabasePrivilegeCheckerTest {
43
44 @Mock
45 private DataSource dataSource;
46
47 @Mock
48 private Connection connection;
49
50 @Mock
51 private DatabaseMetaData metaData;
52
53 @Mock
54 private PreparedStatement preparedStatement;
55
56 @Mock
57 private ResultSet resultSet;
58
59 @BeforeEach
60 void setUp() throws SQLException {
61 when(dataSource.getConnection()).thenReturn(connection);
62 when(connection.getMetaData()).thenReturn(metaData);
63 when(metaData.getUserName()).thenReturn("postgres");
64 when(connection.prepareStatement(anyString())).thenReturn(preparedStatement);
65 when(preparedStatement.executeQuery()).thenReturn(resultSet);
66 when(resultSet.next()).thenReturn(true, false);
67 }
68
69 @Test
70 void assertCheckRolReplication() throws SQLException {
71 PostgreSQLDatabasePrivilegeChecker dataSourceChecker = new PostgreSQLDatabasePrivilegeChecker();
72 when(resultSet.getString("rolreplication")).thenReturn("t");
73 when(resultSet.getString("rolsuper")).thenReturn("f");
74 dataSourceChecker.check(dataSource, PrivilegeCheckType.PIPELINE);
75 verify(resultSet, atLeastOnce()).getString("rolsuper");
76 }
77
78 @Test
79 void assertCheckRolSuper() throws SQLException {
80 PostgreSQLDatabasePrivilegeChecker dataSourceChecker = new PostgreSQLDatabasePrivilegeChecker();
81 when(resultSet.getString("rolsuper")).thenReturn("t");
82 when(resultSet.getString("rolreplication")).thenReturn("f");
83 dataSourceChecker.check(dataSource, PrivilegeCheckType.PIPELINE);
84 verify(resultSet, atLeastOnce()).getString("rolreplication");
85 }
86
87 @Test
88 void assertCheckNoPrivilege() throws SQLException {
89 PostgreSQLDatabasePrivilegeChecker dataSourceChecker = new PostgreSQLDatabasePrivilegeChecker();
90 when(resultSet.getString("rolsuper")).thenReturn("f");
91 when(resultSet.getString("rolreplication")).thenReturn("f");
92 assertThrows(MissingRequiredPrivilegeException.class, () -> dataSourceChecker.check(dataSource, PrivilegeCheckType.PIPELINE));
93 verify(resultSet, atLeastOnce()).getString("rolreplication");
94 }
95 }