from django.test import TestCase

from shop.models import HULL, Shop, ShopSource
from shop.sources import source_assets, source_qty
from shop.tests._corptools import make_asset, make_owner_character

STATION = 60003760
CONTAINER = 900000001


class SourceAssetsTests(TestCase):
    def setUp(self):
        self.user, self.audit = make_owner_character()
        self.shop = Shop.objects.create(slug="s", name="S", created_by=self.user)

    def _src(self, container_item_id=None):
        return ShopSource.objects.create(
            shop=self.shop,
            section=HULL,
            character_audit=self.audit,
            location_id=STATION,
            container_item_id=container_item_id,
        )

    def test_hangar_source_returns_top_level_hangar_assets(self):
        make_asset(self.audit, type_id=671, item_id=1, location_id=STATION)
        make_asset(self.audit, type_id=671, item_id=2, location_id=CONTAINER)
        src = self._src()
        ids = set(source_assets(src).values_list("item_id", flat=True))
        self.assertEqual(ids, {1})

    def test_container_source_returns_container_contents(self):
        make_asset(
            self.audit,
            type_id=671,
            item_id=2,
            location_id=CONTAINER,
            location_flag="Unlocked",
        )
        src = self._src(container_item_id=CONTAINER)
        ids = set(source_assets(src).values_list("item_id", flat=True))
        self.assertEqual(ids, {2})

    def test_source_qty_sums_quantity_by_type(self):
        make_asset(self.audit, type_id=671, item_id=1, location_id=STATION, quantity=3)
        make_asset(self.audit, type_id=671, item_id=2, location_id=STATION, quantity=4)
        make_asset(self.audit, type_id=999, item_id=3, location_id=STATION, quantity=9)

        class FakeType:
            id = 671

        self.assertEqual(source_qty(self._src(), FakeType()), 7)
