From 180a4cd7108cc5f6fa318f59c8d52871287b782f Mon Sep 17 00:00:00 2001
From: "Zhenye (Nathan) Na" <32248549+Zhenye-Na@users.noreply.github.com>
Date: Mon, 23 Oct 2023 05:42:28 +0000
Subject: [PATCH] Add test cases for
 SelectStatementContext.findColumnProjection()

---
 .../statement/SelectStatementBinderTest.java  | 71 +++++++++++++++++++
 1 file changed, 71 insertions(+)

diff --git a/infra/binder/src/test/java/org/apache/shardingsphere/infra/binder/statement/SelectStatementBinderTest.java b/infra/binder/src/test/java/org/apache/shardingsphere/infra/binder/statement/SelectStatementBinderTest.java
index 4e2cc0242aae3..a9603932b8938 100644
--- a/infra/binder/src/test/java/org/apache/shardingsphere/infra/binder/statement/SelectStatementBinderTest.java
+++ b/infra/binder/src/test/java/org/apache/shardingsphere/infra/binder/statement/SelectStatementBinderTest.java
@@ -17,6 +17,11 @@
 
 package org.apache.shardingsphere.infra.binder.statement;
 
+import org.apache.shardingsphere.infra.binder.context.segment.select.projection.Projection;
+import org.apache.shardingsphere.infra.binder.context.segment.select.projection.impl.ColumnProjection;
+import org.apache.shardingsphere.infra.binder.context.segment.select.projection.impl.ExpressionProjection;
+import org.apache.shardingsphere.infra.binder.context.segment.select.projection.impl.SubqueryProjection;
+import org.apache.shardingsphere.infra.binder.context.statement.dml.SelectStatementContext;
 import org.apache.shardingsphere.infra.binder.statement.dml.SelectStatementBinder;
 import org.apache.shardingsphere.infra.database.core.DefaultDatabase;
 import org.apache.shardingsphere.infra.metadata.ShardingSphereMetaData;
@@ -40,19 +45,85 @@
 import java.sql.Types;
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collections;
 import java.util.List;
+import java.util.Optional;
 
 import static org.hamcrest.CoreMatchers.instanceOf;
 import static org.hamcrest.CoreMatchers.is;
 import static org.hamcrest.CoreMatchers.not;
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.mockito.Mockito.RETURNS_DEEP_STUBS;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
 class SelectStatementBinderTest {
     
+    @Test
+    void assertFindColumnProjectionWithoutExpandProjections() {
+        SelectStatementContext selectStatementContext = mock(SelectStatementContext.class, RETURNS_DEEP_STUBS);
+        
+        when(selectStatementContext.getProjectionsContext().getExpandProjections()).thenReturn(Collections.emptyList());
+        
+        Optional<ColumnProjection> result = selectStatementContext.findColumnProjection(1);
+        
+        assertFalse(result.isPresent());
+    }
+    
+    @Test
+    void assertFindColumnProjectionInvalidColumnIndex() {
+        Projection projection = mock(Projection.class);
+        SelectStatementContext selectStatementContext = mock(SelectStatementContext.class, RETURNS_DEEP_STUBS);
+        
+        when(selectStatementContext.getProjectionsContext().getExpandProjections()).thenReturn(Collections.singletonList(projection));
+        
+        Optional<ColumnProjection> result = selectStatementContext.findColumnProjection(2);
+        
+        assertFalse(result.isPresent());
+    }
+    
+    @Test
+    void assertFindColumnProjectionInvalidProjectionType() {
+        ExpressionProjection projection = mock(ExpressionProjection.class);
+        SelectStatementContext selectStatementContext = mock(SelectStatementContext.class, RETURNS_DEEP_STUBS);
+        
+        when(selectStatementContext.getProjectionsContext().getExpandProjections()).thenReturn(Collections.singletonList(projection));
+        
+        Optional<ColumnProjection> result = selectStatementContext.findColumnProjection(1);
+        
+        assertFalse(result.isPresent());
+    }
+    
+    @Test
+    void assertFindColumnProjectionFromColumnProjection() {
+        ColumnProjection projection = mock(ColumnProjection.class);
+        SelectStatementContext selectStatementContext = mock(SelectStatementContext.class, RETURNS_DEEP_STUBS);
+        
+        when(selectStatementContext.getProjectionsContext().getExpandProjections()).thenReturn(Collections.singletonList(projection));
+        
+        Optional<ColumnProjection> result = selectStatementContext.findColumnProjection(1);
+        
+        assertTrue(result.isPresent());
+        assertThat(result.get(), is(projection));
+    }
+    
+    @Test
+    void assertFindColumnProjectionFromSubqueryProjection() {
+        SubqueryProjection projection = mock(SubqueryProjection.class, RETURNS_DEEP_STUBS);
+        ColumnProjection columnProjection = mock(ColumnProjection.class);
+        SelectStatementContext selectStatementContext = mock(SelectStatementContext.class, RETURNS_DEEP_STUBS);
+        
+        when(projection.getProjection()).thenReturn(columnProjection);
+        when(selectStatementContext.getProjectionsContext().getExpandProjections()).thenReturn(Collections.singletonList(projection));
+        
+        Optional<ColumnProjection> result = selectStatementContext.findColumnProjection(1);
+        
+        assertTrue(result.isPresent());
+        assertThat(result.get(), is(columnProjection));
+    }
+    
     @Test
     void assertBind() {
         SelectStatement selectStatement = new MySQLSelectStatement();