import unittest

from connectome.model import Network
from connectome.noise import Subsampling


class TestSubsampling(unittest.TestCase):
    def test_remove_half_of_the_neurons(self):
        net = Network(nr_exc=70, nr_inh=30)
        subsampling = Subsampling(subsampling_fraction=0.5)
        net = subsampling(network=net)
        self.assertEqual(50, net.nr_neurons)

    def test_remove_no_neurons(self):
        net = Network(nr_exc=70, nr_inh=30)
        subsampling = Subsampling(subsampling_fraction=1)
        net = subsampling(network=net)
        self.assertEqual(100, net.nr_neurons)

    def test_remove_all_neurons(self):
        net = Network(nr_exc=70, nr_inh=30)
        subsampling = Subsampling(subsampling_fraction=0)
        net = subsampling(network=net)
        self.assertEqual(0, net.nr_neurons)


if __name__ == "__main__":
    unittest.main()