View Javadoc
1   /*
2    * Licensed to the Apache Software Foundation (ASF) under one or more
3    * contributor license agreements.  See the NOTICE file distributed with
4    * this work for additional information regarding copyright ownership.
5    * The ASF licenses this file to You under the Apache License, Version 2.0
6    * (the "License"); you may not use this file except in compliance with
7    * the License.  You may obtain a copy of the License at
8    *
9    *     http://www.apache.org/licenses/LICENSE-2.0
10   *
11   * Unless required by applicable law or agreed to in writing, software
12   * distributed under the License is distributed on an "AS IS" BASIS,
13   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14   * See the License for the specific language governing permissions and
15   * limitations under the License.
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  }