from unittest import TestCase

from sohstationviewer.model.mseed_data.mseed_helper import (
    retrieve_nets_from_data_dict
)


class TestRetrieveNetsFromDataDict(TestCase):
    def setUp(self):
        self.nets_by_sta = {}
        self.data_dict = {
            'STA1': {'CHA1': {'nets': {'NET1', 'NET2'}},
                     'CHA2': {'nets': {'NET2', 'NET3'}}
                     },
            'STA2': {'CHA1': {'nets': {'NET1'}},
                     'CHA2': {'nets': {'NET1'}}
                     }
            }

    def test_retrieve_nets(self):
        retrieve_nets_from_data_dict(self.data_dict, self.nets_by_sta)
        self.assertEqual(list(self.nets_by_sta.keys()), ['STA1', 'STA2'])
        self.assertEqual(sorted(list(self.nets_by_sta['STA1'])),
                         ['NET1', 'NET2', 'NET3'])
        self.assertEqual(sorted(list(self.nets_by_sta['STA2'])), ['NET1'])