diff --git a/pokemongo_bot/plugin_loader.py b/pokemongo_bot/plugin_loader.py index f3c1fd9c2f..3bded030b3 100644 --- a/pokemongo_bot/plugin_loader.py +++ b/pokemongo_bot/plugin_loader.py @@ -5,11 +5,26 @@ class PluginLoader(object): folder_cache = [] + def _get_correct_path(self, path): + extension = os.path.splitext(path)[1] + + if extension == '.zip': + correct_path = path + else: + correct_path = os.path.dirname(path) + + return correct_path + def load_path(self, path): - parent_dir = os.path.dirname(path) - if parent_dir not in self.folder_cache: - self.folder_cache.append(parent_dir) - sys.path.append(parent_dir) + correct_path = self._get_correct_path(path) + + if correct_path not in self.folder_cache: + self.folder_cache.append(correct_path) + sys.path.append(correct_path) + + def remove_path(self, path): + correct_path = self._get_correct_path(path) + sys.path.remove(correct_path) def get_class(self, namespace_class): [namespace, class_name] = namespace_class.split('.') diff --git a/pokemongo_bot/test/plugin_loader_test.py b/pokemongo_bot/test/plugin_loader_test.py index 2960c5d36b..4d4d5ca952 100644 --- a/pokemongo_bot/test/plugin_loader_test.py +++ b/pokemongo_bot/test/plugin_loader_test.py @@ -18,3 +18,11 @@ def test_load_namespace_class(self): self.plugin_loader.load_path(package_path) loaded_class = self.plugin_loader.get_class('plugin_fixture.FakeTask') self.assertEqual(loaded_class({}, {}).work(), 'FakeTask') + self.plugin_loader.remove_path(package_path) + + def test_load_zip(self): + package_path = os.path.join(os.path.abspath(os.path.dirname(__file__)), 'resources', 'plugin_fixture_test.zip') + self.plugin_loader.load_path(package_path) + loaded_class = self.plugin_loader.get_class('plugin_fixture_test.FakeTask') + self.assertEqual(loaded_class({}, {}).work(), 'FakeTask') + self.plugin_loader.remove_path(package_path) diff --git a/pokemongo_bot/test/resources/plugin_fixture_test.zip b/pokemongo_bot/test/resources/plugin_fixture_test.zip new file mode 100644 index 0000000000..335d95e522 Binary files /dev/null and b/pokemongo_bot/test/resources/plugin_fixture_test.zip differ