fix datastore tests

GoogleCloudPlatform · Dec 13, 2024 · 195fca8 · 195fca8
1 parent 57a9abf
commit 195fca8
Show file tree

Hide file tree

Showing 2 changed files with 56 additions and 21 deletions.
diff --git a/...oogle/cloud/spring/data/datastore/repository/query/DatastoreQueryLookupStrategyTests.java b/...oogle/cloud/spring/data/datastore/repository/query/DatastoreQueryLookupStrategyTests.java
@@ -37,6 +37,7 @@
 import org.springframework.data.repository.query.Parameter;
 import org.springframework.data.repository.query.Parameters;
 import org.springframework.data.repository.query.QueryMethodEvaluationContextProvider;
+import org.springframework.data.repository.query.QueryMethodValueEvaluationContextAccessor;
 import org.springframework.data.repository.query.ValueExpressionDelegate;
 
 /** Tests for the Query Method lookup class. */
@@ -86,6 +87,7 @@ void resolveSqlQueryTest() {
 
     when(namedQueries.hasQuery(queryName)).thenReturn(true);
     when(namedQueries.getQuery(queryName)).thenReturn(query);
+    when(valueExpressionDelegate.getEvaluationContextAccessor()).thenReturn(mock(QueryMethodValueEvaluationContextAccessor.class));
 
     this.datastoreQueryLookupStrategy.resolveQuery(null, null, null, namedQueries);
 

diff --git a/.../java/com/google/cloud/spring/data/datastore/repository/query/GqlDatastoreQueryTests.java b/.../java/com/google/cloud/spring/data/datastore/repository/query/GqlDatastoreQueryTests.java
@@ -48,22 +48,29 @@
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
+import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.function.Supplier;
 import java.util.stream.Stream;
 import org.assertj.core.data.Offset;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.mockito.Mockito;
+import org.mockito.invocation.InvocationOnMock;
+import org.mockito.stubbing.Answer;
 import org.springframework.data.annotation.Id;
 import org.springframework.data.domain.Page;
 import org.springframework.data.domain.PageRequest;
 import org.springframework.data.domain.Slice;
 import org.springframework.data.domain.Sort;
+import org.springframework.data.expression.ValueEvaluationContext;
+import org.springframework.data.expression.ValueEvaluationContextProvider;
 import org.springframework.data.repository.query.Parameter;
 import org.springframework.data.repository.query.Parameters;
 import org.springframework.data.repository.query.QueryMethodEvaluationContextProvider;
+import org.springframework.data.repository.query.QueryMethodValueEvaluationContextAccessor;
 import org.springframework.data.repository.query.ValueExpressionDelegate;
 import org.springframework.expression.EvaluationContext;
 import org.springframework.expression.spel.support.StandardEvaluationContext;
@@ -98,6 +105,7 @@ void initMocks() {
         .thenReturn(this.datastoreEntityConverter);
     when(this.datastoreEntityConverter.getConversions()).thenReturn(this.readWriteConversions);
     this.valueExpressionDelegate = mock(ValueExpressionDelegate.class);
+    when(valueExpressionDelegate.getEvaluationContextAccessor()).thenReturn(mock(QueryMethodValueEvaluationContextAccessor.class));
   }
 
   private GqlDatastoreQuery<Trade> createQuery(
@@ -167,14 +175,7 @@ void compoundNameConventionTest() {
 
     doReturn(key).when(this.datastoreTemplate).getKey(any());
 
-    EvaluationContext evaluationContext = new StandardEvaluationContext();
-    for (int i = 0; i < paramVals.length; i++) {
-      evaluationContext.setVariable(paramNames[i], paramVals[i]);
-    }
-//    when(this.valueExpressionDelegate.getEvaluationContext(any(), any()))
-//        .thenReturn(evaluationContext);
-//    when(this.valueExpressionDelegate.getEvaluationContext(any(), any(), any()))
-//        .thenReturn(evaluationContext);
+    this.valueExpressionDelegate = ValueExpressionDelegate.create();
 
     GqlDatastoreQuery gqlDatastoreQuery = createQuery(gql, false, false);
 
@@ -517,21 +518,53 @@ private Parameters buildParameters(Object[] params, String[] paramNames) {
 
     Mockito.<Parameters>when(this.queryMethod.getParameters()).thenReturn(parameters);
 
-    when(parameters.getNumberOfParameters()).thenReturn(paramNames.length);
-    when(parameters.getParameter(anyInt()))
-        .thenAnswer(
-            invocation -> {
-              int index = invocation.getArgument(0);
-              Parameter param = mock(Parameter.class);
-              when(param.getName())
-                  .thenReturn(
-                      paramNames[index] == null
-                          ? Optional.empty()
-                          : Optional.of(paramNames[index]));
+    final List<Parameter> parameterList = new ArrayList<>();
+    for (int index = 0; index < params.length; index++) {
+
+      Parameter param = mock(Parameter.class);
+      parameterList.add(param);
+      when(param.getName())
+          .thenReturn(
+              paramNames[index] == null ? Optional.empty() : Optional.of(paramNames[index]));
 
-              Mockito.<Class>when(param.getType()).thenReturn(params[index].getClass());
+      Mockito.<Class>when(param.getType()).thenReturn(params[index].getClass());
+      when(param.isNamedParameter()).thenReturn(true);
+      when(param.getRequiredName()).thenReturn(paramNames[index]);
+      when(param.getIndex()).thenReturn(index);
 
-              return param;
+      when(parameters.getParameter(eq(index))).thenAnswer(invocation -> param);
+    }
+    when(parameters.getNumberOfParameters()).thenReturn(paramNames.length);
+
+    Supplier<Iterator> iteratorSupplier =
+        () -> {
+          Iterator<Parameter> parameterIterator = mock(Iterator.class);
+          if (paramNames.length == 0) {
+            when(parameterIterator.hasNext()).thenReturn(false);
+          } else {
+            // every hasNext() call will return true except for the last one
+            Boolean[] hasNextResponses = new Boolean[paramNames.length];
+            Arrays.fill(hasNextResponses, true);
+            hasNextResponses[hasNextResponses.length - 1] = false;
+            when(parameterIterator.hasNext()).thenReturn(true, hasNextResponses);
+
+            // Then we set up the sequence of returned values for the params
+            Parameter[] restOfParameters = new Parameter[parameterList.size() - 1];
+            parameterList.subList(1, parameterList.size()).toArray(restOfParameters);
+            when(parameterIterator.next()).thenReturn(parameterList.getFirst(), restOfParameters);
+          }
+          return parameterIterator;
+        };
+    // we return a new iterator each time. This is because foreach loops call iterable.iterator()
+    // once and parameters
+    // is iterated over several times.
+    when(parameters.iterator())
+        .thenAnswer(
+            new Answer<Iterator>() {
+              @Override
+              public Iterator answer(InvocationOnMock invocation) throws Throwable {
+                return iteratorSupplier.get();
+              }
             });
     return parameters;
   }