from unittest import TestCase

from gamelib.lab import Lab
from gamelib import research, products


class TestLab(TestCase):
    def setUp(self):
        self.lab = Lab()

    def test_find_new_products(self):
        # Set up the appropriate research.
        if research.Tesla in self.lab.new_research:
            self.lab._gain_science(research.Tesla(1))
        self.lab._get_science(research.Tesla).points += 2
        # Check breakthrough options.
        new_products = self.lab.find_new_products(None)
        self.assertTrue(products.TeslaTank in new_products)
        self.assertTrue(products.DoomsdayVirus not in new_products)

    def test_find_new_research(self):
        # Set up the appropriate research.
        if research.Robotics in self.lab.new_research:
            self.lab._gain_science(research.Robotics(1))
        if research.Rocketry in self.lab.new_research:
            self.lab._gain_science(research.Rocketry(1))
        self.lab._get_science(research.Rocketry).points += 1
        # Check breakthrough options.
        new_research = self.lab.find_new_research()
        self.assertTrue(research.Space in new_research)
        self.assertTrue(research.ArtificialIntelligence not in new_research)