Add more tests for import completion

sebastinas · sebastinas · commit 3f7eb13c7c20 · 2021-01-02T19:24:01.000+01:00
Related to #844
diff --git a/bpython/test/test_importcompletion.py b/bpython/test/test_importcompletion.py
@@ -15,19 +15,45 @@ def setUp(self):
             "zzefg",
             "zzabc.e",
             "zzabc.f",
+            "zzefg.a1",
+            "zzefg.a2"
         ]
 
     def test_simple_completion(self):
         self.assertSetEqual(
             self.module_gatherer.complete(10, "import zza"), {"zzabc", "zzabd"}
         )
 
-    def test_package_completion(self):
+    def test_import_empty(self):
         self.assertSetEqual(
             self.module_gatherer.complete(13, "import zzabc."),
             {"zzabc.e", "zzabc.f"},
         )
 
+    def test_import(self):
+        self.assertSetEqual(
+            self.module_gatherer.complete(14, "import zzefg.a"),
+            {"zzefg.a1", "zzefg.a2"},
+        )
+
+
+    @unittest.expectedFailure
+    def test_import_empty(self):
+        self.assertSetEqual(
+            self.module_gatherer.complete(7, "import "), {"zzabc", "zzabd", "zzefg"}
+        )
+
+    @unittest.expectedFailure
+    def test_from_import_empty(self):
+        self.assertSetEqual(
+            self.module_gatherer.complete(18, "from zzabc import "), {"e", "f"}
+        )
+
+    def test_from_import(self):
+        self.assertSetEqual(
+            self.module_gatherer.complete(19, "from zzefg import a"), {"a1", "a2"}
+        )
+
 
 class TestRealComplete(unittest.TestCase):
     def setUp(self):
@@ -52,6 +78,11 @@ def test_from_package(self):
             self.module_gatherer.complete(17, "from xml import d"), {"dom"}
         )
 
+    def test_from_package(self):
+        self.assertSetEqual(
+            self.module_gatherer.complete(17, "from xml import d"), {"dom"}
+        )
+
 
 class TestAvoidSymbolicLinks(unittest.TestCase):
     def setUp(self):
