test: add a annotate test, and simplify the logic

3 files changed, 32 insertions, 5 deletions

diff --git a/coverage/annotate.py b/coverage/annotate.py
index 9ca1b80a..07ff644d 100644
--- a/coverage/annotate.py
+++ b/coverage/annotate.py
@@ -74,8 +74,7 @@ class AnnotateReporter:
             dest_file = fr.filename + ",cover"
 
         with open(dest_file, 'w', encoding='utf-8') as dest:
-            i = 0
-            j = 0
+            i = j = 0
             covered = True
             source = fr.source()
             for lineno, line in enumerate(source.splitlines(True), start=1):
@@ -89,9 +88,7 @@ class AnnotateReporter:
                     dest.write('  ')
                 elif self.else_re.match(line):
                     # Special logic for lines containing only 'else:'.
-                    if i >= len(statements) and j >= len(missing):
-                        dest.write('! ')
-                    elif i >= len(statements) or j >= len(missing):
+                    if j >= len(missing):
                         dest.write('> ')
                     elif statements[i] == missing[j]:
                         dest.write('! ')
diff --git a/tests/gold/annotate/mae/mae.py,cover b/tests/gold/annotate/mae/mae.py,cover
new file mode 100644
index 00000000..ca086e93
--- /dev/null
+++ b/tests/gold/annotate/mae/mae.py,cover
@@ -0,0 +1,10 @@
+> def f(x):
+>     if x == 1:
+>         print("1")
+>     else:
+>         print("2")
+  
+> if f(1):
+!     print("nope")
+> if f(2):
+!     print("nope")
diff --git a/tests/test_annotate.py b/tests/test_annotate.py
index de6edcd0..f85450c7 100644
--- a/tests/test_annotate.py
+++ b/tests/test_annotate.py
@@ -105,3 +105,23 @@ class AnnotationGoldTest(CoverageTest):
         self.start_import_stop(cov, "white")
         cov.annotate()
         compare(gold_path("annotate/annotate"), ".", "*,cover")
+
+    def test_missing_after_else(self):
+        self.make_file("mae.py", """\
+            def f(x):
+                if x == 1:
+                    print("1")
+                else:
+                    print("2")
+
+            if f(1):
+                print("nope")
+            if f(2):
+                print("nope")
+            """)
+
+        cov = coverage.Coverage()
+        self.start_import_stop(cov, "mae")
+        cov.annotate()
+        assert self.stdout() == "1\n2\n"
+        compare(gold_path("annotate/mae"), ".", "*,cover")