diff --git a/src/cstag/consensus.py b/src/cstag/consensus.py
index df17e67..5fc73eb 100644
--- a/src/cstag/consensus.py
+++ b/src/cstag/consensus.py
@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import re
 from itertools import chain
 from collections import deque, Counter
diff --git a/src/cstag/lengthen.py b/src/cstag/lengthen.py
index eb244b2..430ea38 100644
--- a/src/cstag/lengthen.py
+++ b/src/cstag/lengthen.py
@@ -1,5 +1,6 @@
-import re
+from __future__ import annotations
 
+import re
 from cstag.utils.validator import validate_short_format
 
 
diff --git a/src/cstag/mask.py b/src/cstag/mask.py
index 37c7e8b..31a7408 100644
--- a/src/cstag/mask.py
+++ b/src/cstag/mask.py
@@ -1,9 +1,11 @@
+from __future__ import annotations
+
 import re
 
 from cstag.utils.validator import validate_long_format, validate_threshold
 
 
-def mask(cs_tag: str, cigar: str, qual: str, threshold: int = 10, prefix: bool = False):
+def mask(cs_tag: str, cigar: str, qual: str, threshold: int = 10, prefix: bool = False) -> str:
     """Mask low-quality bases to 'N'
     Args:
         cs_tag (str): cs tag in the **long** format
@@ -15,11 +17,11 @@ def mask(cs_tag: str, cigar: str, qual: str, threshold: int = 10, prefix: bool =
         str: Masked cs tag
     Example:
         >>> import cstag
-        >>> cs_tag = "cs:Z:=ACGT*ac+gg-cc=T"
+        >>> cs_tag = "=ACGT*ac+gg-cc=T"
         >>> cigar = "5M2I2D1M"
         >>> qual = "AA!!!!AA"
         >>> cstag.mask(cs_tag, qual)
-        cs:Z:=ACNN*an+ng-cc=T
+        =ACNN*an+ng-cc=T
     """
 
     validate_long_format(cs_tag)
diff --git a/src/cstag/shorten.py b/src/cstag/shorten.py
index 5783b45..baa20e1 100644
--- a/src/cstag/shorten.py
+++ b/src/cstag/shorten.py
@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import re
 
 
diff --git a/src/cstag/to_html.py b/src/cstag/to_html.py
index e0dc832..48604bb 100644
--- a/src/cstag/to_html.py
+++ b/src/cstag/to_html.py
@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import re
 
 from cstag.utils.validator import validate_long_format
diff --git a/tests/test_mask.py b/tests/test_mask.py
index f1e7dd3..c6a8b34 100644
--- a/tests/test_mask.py
+++ b/tests/test_mask.py
@@ -16,6 +16,13 @@ def test_basic_with_prefix():
     assert cstag.mask(CSTAG, CIGAR, QUAL, prefix=True) == "cs:Z:=ACNN*an+ng-cc=T"
 
 
+def test_all_n():
+    CSTAG = "=ACGT"
+    CIGAR = "4M"
+    QUAL = "!!!!"
+    assert cstag.mask(CSTAG, CIGAR, QUAL) == "=NNNN"
+
+
 def test_softclip():
     CSTAG = "=ACGT*ac+gg-cc=T"
     CIGAR = "2S5M2I2D1M"