from unittest import TestCase

from gamelib.lab import Lab, ScienceGraph
from gamelib import research, schematics


LAB_DATA = {
    'science': {
        'research.Physics': 5,
        'research.Rocketry': 2,
        'research.Electrickery': 3,
        'schematic.MachineGun': 1,
        },
    }


class TestLab(TestCase):
    def test_new_lab(self):
        all_sciences = set()
        for _ in range(10):
            lab = Lab()
            research_areas = [r for r in lab.science
                              if isinstance(r, research.ResearchArea)]
            self.assertEqual(3, len(research_areas))
            for science in research_areas:
                all_sciences.add(type(science))
        self.assertTrue(len(all_sciences) > 3)

    def test_find_new_schematics(self):
        lab = Lab(LAB_DATA)
        new_schematics = lab.find_new_schematics()
        self.assertTrue(schematics.LightningGun in new_schematics)
        self.assertTrue(schematics.DoomsdayVirus not in new_schematics)

    def test_find_new_research(self):
        lab = Lab(LAB_DATA)
        new_research = lab.find_new_research()
        self.assertTrue(research.Lasers in new_research)
        self.assertTrue(research.ArtificialIntelligence not in new_research)

    def test_save_data(self):
        # Check that we save what we loaded.
        lab = Lab(LAB_DATA)
        self.assertEqual(LAB_DATA, lab.save_data())
        # Add a schematic and check that it gets saved as well.
        lab._gain_science(schematics.LightningGun())
        new_science = dict(
            LAB_DATA['science'].items() + [('schematic.LightningGun', 0)])
        self.assertEqual({'science': new_science}, lab.save_data())

    def test_science_graph(self):
        lab = Lab(LAB_DATA)
        self.assertEqual((True, [research.Electrickery, research.Physics]),
                         lab.suggest_research())