Fix matplotlib display error

ljvmiranda921 · May 17, 2018 · a62f17b · a62f17b
1 parent 8d72e11
commit a62f17b
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 3 deletions.
diff --git a/tests/utils/environments/conftest.py b/tests/utils/environments/conftest.py
@@ -4,9 +4,14 @@
 """Fixtures for tests"""
 
 # Import modules
+import os
 import pytest
 import numpy as np
 from mock import Mock
+import matplotlib as mpl
+
+if os.environ.get('DISPLAY','') == '':
+    mpl.use('Agg')
 
 # Import from package
 from pyswarms.single import GlobalBestPSO

diff --git a/tests/utils/environments/test_plot_environment.py b/tests/utils/environments/test_plot_environment.py
@@ -3,13 +3,13 @@
 import pytest
 import numpy as np
 import matplotlib as mpl
-from matplotlib.axes._subplots import SubplotBase
-from matplotlib.animation import FuncAnimation
 
-# Fix for display
 if os.environ.get('DISPLAY','') == '':
     mpl.use('Agg')
 
+from matplotlib.axes._subplots import SubplotBase
+from matplotlib.animation import FuncAnimation
+
 # Import from package
 from pyswarms.utils.environments import PlotEnvironment
 from pyswarms.utils.functions.single_obj import sphere_func