From 85333e08d90863d701196af0152a1dd84000188c Mon Sep 17 00:00:00 2001
From: hanhxiao <hanhxiao@tencent.com>
Date: Sat, 20 Jul 2019 01:00:00 +0800
Subject: [PATCH] feat(base): support loading external modules from py and yaml

---
 tests/test_contrib_module.py          | 35 ++++++++++++++++++++++-----
 tests/test_contrib_module_negative.py | 21 ----------------
 2 files changed, 29 insertions(+), 27 deletions(-)
 delete mode 100644 tests/test_contrib_module_negative.py

diff --git a/tests/test_contrib_module.py b/tests/test_contrib_module.py
index c71596ab..0206e6be 100644
--- a/tests/test_contrib_module.py
+++ b/tests/test_contrib_module.py
@@ -1,19 +1,42 @@
 import os
-import unittest
+import sys
+import unittest.mock
 
+import ruamel.yaml
+
+dirname = os.path.dirname(__file__)
+module_path = os.path.join(dirname, 'contrib', 'dummy_contrib.py')
+cls_name = 'FooContribEncoder'
 
-class TestYaml(unittest.TestCase):
 
+class TestYaml(unittest.TestCase):
     def setUp(self):
-        dirname = os.path.dirname(__file__)
-        module_path = os.path.join(dirname, 'contrib', 'dummy_contrib.py')
-        cls_name = 'FooContribEncoder'
-        os.environ['GNES_CONTRIB_MODULE'] = '%s:%s' % (cls_name, module_path)
         self.yaml_path = os.path.join(os.path.dirname(__file__),
                                       'contrib', 'dummy.yml')
 
+        # reload gnes module on every unit test
+        for mod in list(sys.modules.keys()):
+            if mod.startswith('gnes.'):
+                del (sys.modules[mod])
+
+    @unittest.mock.patch.dict(os.environ, {'GNES_CONTRIB_MODULE': '%s:%s' % (cls_name, module_path)})
     def test_load_contrib(self):
+
         from gnes.encoder.base import BaseEncoder, BaseTextEncoder
         a = BaseEncoder.load_yaml(self.yaml_path)
         self.assertIsInstance(a, BaseTextEncoder)
         self.assertEqual(a.encode([]), 'hello 531')
+
+    @unittest.mock.patch.dict(os.environ, {'GNES_CONTRIB_MODULE': '%s:%s' % ('blah', module_path)})
+    def test_bad_name(self):
+        try:
+            from gnes.encoder.base import BaseEncoder
+        except AttributeError:
+            pass
+
+    @unittest.mock.patch.dict(os.environ, {'GNES_CONTRIB_MODULE': '%s:%s' % (cls_name, 'blah')})
+    def test_bad_path(self):
+        try:
+            from gnes.encoder.base import BaseEncoder
+        except AttributeError:
+            pass
diff --git a/tests/test_contrib_module_negative.py b/tests/test_contrib_module_negative.py
deleted file mode 100644
index 10e1952a..00000000
--- a/tests/test_contrib_module_negative.py
+++ /dev/null
@@ -1,21 +0,0 @@
-import os
-import unittest
-
-import ruamel.yaml
-
-
-class TestYaml(unittest.TestCase):
-
-    def setUp(self):
-        dirname = os.path.dirname(__file__)
-        module_path = os.path.join(dirname, 'contrib', 'dummy_contrib.py')
-        cls_name = 'FooContribEncoder'
-        os.environ['GNES_CONTRIB_MODULE'] = '%s:%s' % (cls_name, module_path)
-        self.yaml_path = os.path.join(os.path.dirname(__file__),
-                                      'contrib', 'dummy.yml')
-
-    def test_broken_contrib(self):
-        os.environ['GNES_CONTRIB_MODULE'] = ''
-        from gnes.encoder.base import BaseEncoder
-
-        self.assertRaises(ruamel.yaml.constructor.ConstructorError, BaseEncoder.load_yaml, self.yaml_path)