Add attribute access test

2 files changed, 29 insertions, 2 deletions

diff --git a/config-model/src/test/java/com/yahoo/searchdefinition/processing/RankProfileSearchFixture.java b/config-model/src/test/java/com/yahoo/searchdefinition/processing/RankProfileSearchFixture.java
index 8ba0a42f799..79bdbddbdd6 100644
--- a/config-model/src/test/java/com/yahoo/searchdefinition/processing/RankProfileSearchFixture.java
+++ b/config-model/src/test/java/com/yahoo/searchdefinition/processing/RankProfileSearchFixture.java
@@ -29,9 +29,14 @@ class RankProfileSearchFixture {
     }
 
     RankProfileSearchFixture(ApplicationPackage applicationpackage, String rankProfiles) throws ParseException {
+        this(applicationpackage, rankProfiles, null);
+    }
+
+    RankProfileSearchFixture(ApplicationPackage applicationpackage, String rankProfiles, String field) throws ParseException {
         SearchBuilder builder = new SearchBuilder(applicationpackage, rankProfileRegistry);
         String sdContent = "search test {\n" +
                            "  document test {\n" +
+                           (field != null ? field : "") +
                            "  }\n" +
                            rankProfiles +
                            "\n" +
diff --git a/config-model/src/test/java/com/yahoo/searchdefinition/processing/RankingExpressionWithTensorFlowTestCase.java b/config-model/src/test/java/com/yahoo/searchdefinition/processing/RankingExpressionWithTensorFlowTestCase.java
index 598ed04e657..60a2c5326a0 100644
--- a/config-model/src/test/java/com/yahoo/searchdefinition/processing/RankingExpressionWithTensorFlowTestCase.java
+++ b/config-model/src/test/java/com/yahoo/searchdefinition/processing/RankingExpressionWithTensorFlowTestCase.java
@@ -69,6 +69,19 @@ public class RankingExpressionWithTensorFlowTestCase {
                                                                               queryProfileType);
         RankProfileSearchFixture search = fixtureWith("query(mytensor)",
                                                       "tensorflow('mnist_softmax/saved')",
+                                                      null,
+                                                      application);
+        search.assertFirstPhaseExpression(vespaExpression, "my_profile");
+        assertConstant("Variable_1", search, Optional.of(10L));
+        assertConstant("Variable", search, Optional.of(7840L));
+    }
+
+    @Test
+    public void testTensorFlowReferenceWithDocumentFeature() throws ParseException {
+        StoringApplicationPackage application = new StoringApplicationPackage(applicationDir);
+        RankProfileSearchFixture search = fixtureWith("attribute(mytensor)",
+                                                      "tensorflow('mnist_softmax/saved')",
+                                                      "field mytensor type tensor(d0[],d1[784]) { indexing: attribute }",
                                                       application);
         search.assertFirstPhaseExpression(vespaExpression, "my_profile");
         assertConstant("Variable_1", search, Optional.of(10L));
@@ -186,6 +199,7 @@ public class RankingExpressionWithTensorFlowTestCase {
             StoringApplicationPackage storedApplication = new StoringApplicationPackage(storedApplicationDirectory);
             RankProfileSearchFixture searchFromStored = fixtureWith("tensor(d0[2],d1[784])(0.0)",
                                                                     "tensorflow('mnist_softmax/saved')",
+                                                                    null,
                                                                     storedApplication);
             searchFromStored.assertFirstPhaseExpression(vespaExpression, "my_profile");
             // Verify that the constants exists, but don't verify the content as we are not
@@ -225,11 +239,18 @@ public class RankingExpressionWithTensorFlowTestCase {
     }
 
     private RankProfileSearchFixture fixtureWith(String placeholderExpression, String firstPhaseExpression) {
-        return fixtureWith(placeholderExpression, firstPhaseExpression, new StoringApplicationPackage(applicationDir));
+        return fixtureWith(placeholderExpression, firstPhaseExpression, null,
+                           new StoringApplicationPackage(applicationDir));
+    }
+
+    private RankProfileSearchFixture fixtureWith(String placeholderExpression, String firstPhaseExpression, String field) {
+        return fixtureWith(placeholderExpression, firstPhaseExpression, field,
+                           new StoringApplicationPackage(applicationDir));
     }
 
     private RankProfileSearchFixture fixtureWith(String placeholderExpression,
                                                  String firstPhaseExpression,
+                                                 String field,
                                                  StoringApplicationPackage application) {
         try {
             return new RankProfileSearchFixture(
@@ -241,7 +262,8 @@ public class RankingExpressionWithTensorFlowTestCase {
                     "    first-phase {\n" +
                     "      expression: " + firstPhaseExpression +
                     "    }\n" +
-                    "  }");
+                    "  }",
+                    field);
         }
         catch (ParseException e) {
             throw new IllegalArgumentException(e);