Correct wrong assumption about possible member attribute positions

Apparently, they always have to be on the first line after the start of the member.
rubberduck-vba · Dec 15, 2018 · f668115 · f668115
1 parent ff6ba2e
commit f668115
Show file tree

Hide file tree

Showing 4 changed files with 65 additions and 11 deletions.
diff --git a/Rubberduck.Parsing/ParserRuleContextExtensions.cs b/Rubberduck.Parsing/ParserRuleContextExtensions.cs
@@ -198,7 +198,7 @@ public static bool ContainsTokenIndex(this ParserRuleContext context, int tokenI
         public static TContext GetDescendent<TContext>(this ParserRuleContext context) where TContext : ParserRuleContext
         {
             var descendents = GetDescendents<TContext>(context);
-            return descendents.FirstOrDefault();
+            return descendents.OrderBy(descendent => descendent.Start.TokenIndex).FirstOrDefault();
         }
 
         /// <summary>

diff --git a/Rubberduck.Parsing/VBA/AttributesUpdater.cs b/Rubberduck.Parsing/VBA/AttributesUpdater.cs
@@ -74,8 +74,18 @@ public void AddAttribute(IRewriteSession rewriteSession, Declaration declaration
             }
             else
             {
-                var codeToInsert = $"{Environment.NewLine}Attribute {attribute} ={AttributeValuesText(values)}";
-                rewriter.InsertAfter(declaration.AttributesPassContext.Stop.TokenIndex, codeToInsert);
+                ParserRuleContext attributesContext = declaration.AttributesPassContext;
+                var firstEndOfLineInMember = attributesContext.GetDescendent<VBAParser.EndOfLineContext>();
+                if (firstEndOfLineInMember == null)
+                {
+                    var codeToInsert = $"{Environment.NewLine}Attribute {attribute} ={AttributeValuesText(values)}";
+                    rewriter.InsertAfter(declaration.AttributesPassContext.Stop.TokenIndex, codeToInsert);
+                }
+                else
+                {
+                    var codeToInsert = $"Attribute {attribute} ={AttributeValuesText(values)}{Environment.NewLine}";
+                    rewriter.InsertAfter(firstEndOfLineInMember.Stop.TokenIndex, codeToInsert);
+                }
             }
         }
 

diff --git a/RubberduckTests/PostProcessing/AttributesUpdaterTests.cs b/RubberduckTests/PostProcessing/AttributesUpdaterTests.cs
@@ -14,7 +14,7 @@ public class AttributesUpdaterTests
     {
         [Test]
         [Category("AttributesUpdater")]
-        public void AddAttributeAddsMemberAttributeBelowMember()
+        public void AddAttributeAddsMemberAttributeBelowFirstLineOfMember()
         {
             const string inputCode =
                 @"VERSION 1.0 CLASS
@@ -36,9 +36,9 @@ End Sub
 Attribute VB_Name = ""ClassKeys""
 Attribute VB_GlobalNameSpace = False
 Public Sub Foo(bar As String)
+Attribute Foo.VB_Description = ""The MyFunc Description""
     bar = vbNullString
 End Sub
-Attribute Foo.VB_Description = ""The MyFunc Description""
 ";
             var attributeToAdd = "Foo.VB_Description";
             var attributeValues = new List<string> {"\"The MyFunc Description\""};
@@ -60,6 +60,50 @@ End Sub
             Assert.AreEqual(expectedCode, actualCode);
         }
 
+        [Test]
+        [Category("AttributesUpdater")]
+        public void AddAttributeAddsMemberAttributeBelowOneLineMember()
+        {
+            const string inputCode =
+                @"VERSION 1.0 CLASS
+BEGIN
+  MultiUse = -1  'True
+END
+Attribute VB_Name = ""ClassKeys""
+Attribute VB_GlobalNameSpace = False
+Public Sub Foo(bar As String) : bar = vbNullString : End Sub
+";
+
+            const string expectedCode =
+                @"VERSION 1.0 CLASS
+BEGIN
+  MultiUse = -1  'True
+END
+Attribute VB_Name = ""ClassKeys""
+Attribute VB_GlobalNameSpace = False
+Public Sub Foo(bar As String) : bar = vbNullString : End Sub
+Attribute Foo.VB_Description = ""The MyFunc Description""
+";
+            var attributeToAdd = "Foo.VB_Description";
+            var attributeValues = new List<string> { "\"The MyFunc Description\"" };
+
+            string actualCode;
+            var (component, rewriteSession, state) = TestSetup(inputCode);
+            using (state)
+            {
+                var fooDeclaration = state.DeclarationFinder
+                    .UserDeclarations(DeclarationType.Procedure)
+                    .First(decl => decl.IdentifierName == "Foo");
+                var attributesUpdater = new AttributesUpdater(state);
+
+                attributesUpdater.AddAttribute(rewriteSession, fooDeclaration, attributeToAdd, attributeValues);
+                rewriteSession.TryRewrite();
+
+                actualCode = component.CodeModule.Content();
+            }
+            Assert.AreEqual(expectedCode, actualCode);
+        }
+
         [Test]
         [Category("AttributesUpdater")]
         public void MultipleAddAttributeWorkForMembers()
@@ -84,10 +128,10 @@ End Sub
 Attribute VB_Name = ""ClassKeys""
 Attribute VB_GlobalNameSpace = False
 Public Sub Foo(bar As String)
-    bar = vbNullString
-End Sub
 Attribute Foo.VB_Description = ""The MyFunc Description""
 Attribute Foo.VB_HelpID = 2
+    bar = vbNullString
+End Sub
 ";
             var firstAttributeToAdd = "Foo.VB_Description";
             var firstAttributeValues = new List<string> { "\"The MyFunc Description\"" };

diff --git a/RubberduckTests/QuickFixes/AddMissingAttributeQuickFixTests.cs b/RubberduckTests/QuickFixes/AddMissingAttributeQuickFixTests.cs
@@ -44,9 +44,9 @@ public void MissingMemberAttribute_QuickFixWorks()
             const string expectedCode =
                 @"'@MemberAttribute VB_Description, ""Desc""
 Public Sub Foo()
+Attribute Foo.VB_Description = ""Desc""
     Const const1 As Integer = 9
-End Sub
-Attribute Foo.VB_Description = ""Desc""";
+End Sub";
 
             var actualCode = ApplyQuickFixToFirstInspectionResult(inputCode, state => new MissingAttributeInspection(state), CodeKind.AttributesCode);
             Assert.AreEqual(expectedCode, actualCode);
@@ -86,9 +86,9 @@ public void MissingMemberAttributeWithMultipleValues_QuickFixWorks()
             const string expectedCode =
                 @"'@MemberAttribute VB_Ext_Key, ""Key"", ""Value""
 Public Sub Foo()
+Attribute Foo.VB_Ext_Key = ""Key"", ""Value""
     Const const1 As Integer = 9
-End Sub
-Attribute Foo.VB_Ext_Key = ""Key"", ""Value""";
+End Sub";
 
             var actualCode = ApplyQuickFixToFirstInspectionResult(inputCode, state => new MissingAttributeInspection(state), CodeKind.AttributesCode);
             Assert.AreEqual(expectedCode, actualCode);