fix(ast): add import generation for annotations on VariableExpr (#1076)

* This PR is a follow-up patch to #1012 which added support for annotations on VariableExpr. It updates ImportWriterVisitor so that import generation also covers these annotations.
googleapis · Nov 3, 2022 · f5d5524 · f5d5524
1 parent 7941431
commit f5d5524
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 0 deletions.
diff --git a/src/main/java/com/google/api/generator/engine/writer/ImportWriterVisitor.java b/src/main/java/com/google/api/generator/engine/writer/ImportWriterVisitor.java
@@ -183,6 +183,7 @@ public void visit(TernaryExpr ternaryExpr) {
   @Override
   public void visit(VariableExpr variableExpr) {
     variableExpr.variable().type().accept(this);
+    annotations(variableExpr.annotations());
     if (variableExpr.exprReferenceExpr() != null) {
       variableExpr.exprReferenceExpr().accept(this);
     }

diff --git a/src/test/java/com/google/api/generator/engine/writer/ImportWriterVisitorTest.java b/src/test/java/com/google/api/generator/engine/writer/ImportWriterVisitorTest.java
@@ -363,6 +363,32 @@ public void writeVariableExprImports_nestedReference() {
         writerVisitor.write());
   }
 
+  @Test
+  public void writeVariableExprImports_withAnnotations() {
+    Variable variable =
+        Variable.builder()
+            .setName("expr")
+            .setType(TypeNode.withReference(ConcreteReference.withClazz(Expr.class)))
+            .build();
+
+    VariableExpr variableExpr =
+        VariableExpr.builder()
+            .setVariable(variable)
+            .setIsDecl(true)
+            .setAnnotations(
+                Arrays.asList(
+                    AnnotationNode.withType(
+                        TypeNode.withReference(ConcreteReference.withClazz(Generated.class)))))
+            .build();
+
+    variableExpr.accept(writerVisitor);
+    assertEquals(
+        LineFormatter.lines(
+            "import com.google.api.generator.engine.ast.Expr;\n",
+            "import javax.annotation.Generated;\n\n"),
+        writerVisitor.write());
+  }
+
   @Test
   public void writeAnonymousClassExprImports() {
     // [Constructing] Function<List<IOException>, MethodDefinition>