Simulator Bug Fixes (#143)

* Convert assertion to exception and add corresponding test

* Adds function to fetch all channels

stingray/simulator/simulator.py

@@ -170,8 +170,11 @@ def simulate_channel(self, channel, *args):
         """
 
         # Check that channel name does not already exist.
-        assert channel not in [lc[0] for lc in self.channels]
-        self.channels.append((channel, self.simulate(*args)))
+        if channel not in [lc[0] for lc in self.channels]:
+            self.channels.append((channel, self.simulate(*args)))
+
+        else:
+            raise KeyError('A channel with this name already exists.')
 
     def get_channel(self, channel):
         """
@@ -187,6 +190,13 @@ def get_channels(self, channels):
 
         return [lc[1] for lc in self.channels if lc[0] in channels]
 
+    def get_all_channels(self):
+        """
+        Get lightcurves belonging to all channels.
+        """
+
+        return [lc[1] for lc in self.channels]
+
     def delete_channel(self, channel):
         """
         Delete an energy channel.

stingray/simulator/tests/test_simulator.py

@@ -53,6 +53,13 @@ def test_simulate_channel(self):
         self.simulator.simulate_channel('3.5-4.5', 'lorenzian', [1, 2, 3, 4])
         self.simulator.delete_channel('3.5-4.5')
 
+    def test_incorrect_simulate_channel(self):
+        """Test simulating a channel that already exists."""
+        self.simulator.simulate_channel('3.5-4.5', 2)
+        with pytest.raises(KeyError):
+            self.simulator.simulate_channel('3.5-4.5', 2)
+        self.simulator.delete_channel('3.5-4.5')
+
     def test_get_channel(self):
         """
         Retrieve an energy channel after it has been simulated.
@@ -71,6 +78,14 @@ def test_get_channels(self):
 
         self.simulator.delete_channels(['3.5-4.5', '4.5-5.5'])
 
+    def test_get_all_channels(self):
+        """ Retrieve all energy channels. """
+        self.simulator.simulate_channel('3.5-4.5', 2)
+        self.simulator.simulate_channel('4.5-5.5', 1)
+        lc = self.simulator.get_all_channels()
+
+        self.simulator.delete_channels(['3.5-4.5', '4.5-5.5'])
+
     def test_count_channels(self):
         """
         Count energy channels after they have been simulated.