Add testSessionRepresentationReturnsCorrectGroupsDuringImpersonation

Author: asu80

testing/trino-tests/src/test/java/io/trino/security/TestImpersonation.java

@@ -15,10 +15,12 @@
 
 import com.google.common.collect.ImmutableSet;
 import io.airlift.log.Logging;
+import io.trino.Session;
 import io.trino.jdbc.TrinoConnection;
 import io.trino.metadata.SystemSecurityMetadata;
 import io.trino.plugin.memory.MemoryPlugin;
 import io.trino.server.testing.TestingTrinoServer;
+import io.trino.spi.security.Identity;
 import io.trino.spi.security.TrinoPrincipal;
 import io.trino.testing.TestingAccessControlManager;
 import org.junit.jupiter.api.BeforeAll;
@@ -34,11 +36,13 @@
 import java.sql.SQLException;
 import java.sql.Statement;
 import java.util.Optional;
+import java.util.Set;
 import java.util.stream.Stream;
 
 import static com.google.inject.multibindings.OptionalBinder.newOptionalBinder;
 import static io.trino.jdbc.BaseTrinoDriverTest.getCurrentUser;
 import static io.trino.spi.security.PrincipalType.USER;
+import static io.trino.testing.TestingSession.testSessionBuilder;
 import static java.lang.String.format;
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
@@ -138,6 +142,31 @@ public void testImpersonationDisallowedWhenRoleIsNone()
         }
     }
 
+    @Test
+    @Timeout(10)
+    public void testSessionRepresentationReturnsCorrectGroupsDuringImpersonation()
+    {
+        // given
+        Set<String> aliceGroups = ImmutableSet.of("alice_group");
+        Set<String> johnGroups = ImmutableSet.of("john_group");
+        Identity alice = Identity.forUser("alice").withGroups(aliceGroups).build();
+        Identity john = Identity.forUser("john").withGroups(johnGroups).build();
+
+        // when
+        Session aliceImpersonationSession = testSessionBuilder()
+                .setOriginalIdentity(alice)
+                .setIdentity(john)
+                .build();
+
+        // then
+        Set<String> originalUserGroups = aliceImpersonationSession.toSessionRepresentation()
+                .getOriginalUserGroups();
+        Set<String> userGroups = aliceImpersonationSession.toSessionRepresentation()
+                .getGroups();
+        assertThat(originalUserGroups).isEqualTo(aliceGroups);
+        assertThat(userGroups).isEqualTo(johnGroups);
+    }
+
     private Connection createConnection(String catalog, String schema, String user)
             throws SQLException
     {