From 080ccac047fb4fea7ed4a8d201cea5597636fe4c Mon Sep 17 00:00:00 2001
From: Hantao Cui <cuihantao@gmail.com>
Date: Thu, 28 Oct 2021 23:09:00 -0500
Subject: [PATCH] Style fixes.

---
 andes/system.py         |  6 +++---
 andes/utils/__init__.py |  2 +-
 tests/test_snapshot.py  | 15 +++++++++++++--
 3 files changed, 17 insertions(+), 6 deletions(-)

diff --git a/andes/system.py b/andes/system.py
index 6bb314b62..0d53434fb 100644
--- a/andes/system.py
+++ b/andes/system.py
@@ -17,9 +17,10 @@
 import os
 import sys
 import inspect
-import dill
 import time
 
+import dill
+
 from collections import OrderedDict
 from typing import Dict, Tuple, Union, Optional
 
@@ -37,9 +38,8 @@
 from andes.io.streaming import Streaming
 
 from andes.shared import np, jac_names, dilled_vars
-from andes.shared import matrix, spmatrix, sparse, Pool
+from andes.shared import matrix, spmatrix, sparse, Pool, Process
 
-from multiprocess import Process
 
 logger = logging.getLogger(__name__)
 dill.settings['recurse'] = True
diff --git a/andes/utils/__init__.py b/andes/utils/__init__.py
index fcc7b1bc0..a94fdfa65 100644
--- a/andes/utils/__init__.py
+++ b/andes/utils/__init__.py
@@ -1,4 +1,4 @@
 from andes.utils import paths           # NOQA
 from andes.utils import sympy           # NOQA
 from andes.utils.paths import get_case  # NOQA
-from andes.utils import snapshot        # NOQA
\ No newline at end of file
+from andes.utils import snapshot        # NOQA
diff --git a/tests/test_snapshot.py b/tests/test_snapshot.py
index 1ff01c2ef..5941f43dd 100644
--- a/tests/test_snapshot.py
+++ b/tests/test_snapshot.py
@@ -1,8 +1,12 @@
+"""
+Test ANDES snapshot based on dill.
+"""
+
 import os
-import andes
+import unittest
 import numpy as np
 
-import unittest
+import andes
 from andes.utils.snapshot import save_ss, load_ss
 
 
@@ -12,6 +16,10 @@ class TestSnapshot(unittest.TestCase):
     """
 
     def test_save_ss(self):
+        """
+        Test saving a snapshot.
+        """
+
         ss = andes.run(andes.get_case("kundur/kundur_full.xlsx"))
         ss.TDS.config.tf = 2
         ss.TDS.run()
@@ -20,6 +28,9 @@ def test_save_ss(self):
         os.remove('test_ss.pkl')
 
     def load_ss(self):
+        """
+        Test loading a snapshot and continuing the simulation.
+        """
 
         # load a snapshot
         test_dir = os.path.dirname(__file__)