Improve get_future_imports implementation.

[etc/vim.git] / tests / test_black.py

diff --git a/tests/test_black.py b/tests/test_black.py
index 8a371973b28c1a4e1192ec41846778e8f793ac2e..cc53aa61221e346fc144b2106a2a3077aca9f193 100644 (file)
--- a/tests/test_black.py
+++ b/tests/test_black.py
@@ -735,6 +735,14 @@ class BlackTestCase(unittest.TestCase):
         self.assertEqual(set(), black.get_future_imports(node))
         node = black.lib2to3_parse("from some.module import black\n")
         self.assertEqual(set(), black.get_future_imports(node))
+        node = black.lib2to3_parse(
+            "from __future__ import unicode_literals as _unicode_literals"
+        )
+        self.assertEqual({"unicode_literals"}, black.get_future_imports(node))
+        node = black.lib2to3_parse(
+            "from __future__ import unicode_literals as _lol, print"
+        )
+        self.assertEqual({"unicode_literals", "print"}, black.get_future_imports(node))
 
     def test_debug_visitor(self) -> None:
         source, _ = read_data("debug_visitor.py")