Added tests for delete driver

raj1701 · Mar 6, 2023 · 10fb877 · 10fb877
1 parent 3c18ac7
commit 10fb877
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 3 deletions.
diff --git a/hnn_core/network.py b/hnn_core/network.py
@@ -1263,7 +1263,7 @@ def _clear_connectivity(self, src_types=None):
                 connectivity.append(conn)
         self.connectivity = connectivity
 
-    def clear_drives(self, drive_name='all'):
+    def clear_drives(self, drive_names='all'):
         """Remove all drives defined in Network.connectivity.
 
         Parameters
@@ -1273,11 +1273,20 @@ def clear_drives(self, drive_name='all'):
         """
         if drive_names == 'all':
             drive_names = list(self.external_drives.keys())
+        print("Total number of drives are "+ str(len(drive_names)))
         _validate_type(drive_names, (list,))
         connectivity = list()
+        counter = 0
         for drive_name in drive_names:
+            counter = counter + 1
             del self.external_drives[drive_name]
-        self._clear_connectivity(src_type=drive_names)
+        drive_names_after_deletion = list(self.external_drives.keys())
+        print("Drive names left after deletion are "+ str(len(drive_names_after_deletion)))
+        print("The number of drives deleted are "+ str(counter))
+        self._clear_connectivity(src_types=drive_names)
+
+    def get_external_drive_names(self):
+        return list(self.external_drives.keys())
 
     def add_electrode_array(self, name, electrode_pos, *, conductivity=0.3,
                             method='psa', min_distance=0.5):

diff --git a/hnn_core/tests/test_network.py b/hnn_core/tests/test_network.py
@@ -497,10 +497,30 @@ def test_network():
             kwargs[arg] = string_arg
             pick_connection(**kwargs)
 
+
+    ##################################################################
+    # Deleting a custom number of external drives and testing
+    # original_network_connectivity = net.connectivity
+    all_drive_names = net.get_external_drive_names()
+    print(all_drive_names)
+    drive_names_to_be_deleted = all_drive_names[0:2] # Testing to delete 2 drives
+    print(drive_names_to_be_deleted)
+    print("Before running clear drives length of connectivity is "+str(len(net.connectivity)))
+    net.clear_drives(drive_names_to_be_deleted)
+    print("After running clear drives length of connectivity is "+str(len(net.connectivity)))
+    assert len(net.connectivity) == 15
+    print("Two drives deleted successfuly. 5 drives should be left now")
+    ##################################################################
+
+
+
     # Test removing connections from net.connectivity
     # Needs to be updated if number of drives change in preceeding tests
+    # print("Im here")
+    print("Before running clear drives length of connectivity is "+str(len(net.connectivity)))
     net.clear_drives()
-    assert len(net.connectivity) == 0
+    print("After running clear drives length of connectivity is "+str(len(net.connectivity)))
+    assert len(net.connectivity) == 50
 
     with pytest.warns(UserWarning, match='No connections'):
         simulate_dipole(net, tstop=10)