from django.test import TestCase
from corehq.apps.domain.shortcuts import create_domain
from corehq.apps.locations.models import LocationType


class TestLocationTypes(TestCase):
    def setUp(self):
        self.domain = create_domain('locations-test')

        def make_loc_type(name, parent_type=None):
            return LocationType.objects.create(
                domain=self.domain.name,
                name=name,
                code=name,
                parent_type=parent_type,
            )

        self.state = make_loc_type('state')

        self.district = make_loc_type('district', self.state)
        self.section = make_loc_type('section', self.district)
        self.block = make_loc_type('block', self.district)
        self.center = make_loc_type('center', self.block)

        self.county = make_loc_type('county', self.state)
        self.city = make_loc_type('city', self.county)

    def tearDown(self):
        self.domain.delete()
        LocationType.objects.filter(domain=self.domain.name).delete()

    def test_hierarchy(self):
        hierarchy = LocationType.objects.full_hierarchy(self.domain.name)
        desired_hierarchy = {
            self.state.id: (
                self.state,
                {
                    self.district.id: (
                        self.district,
                        {
                            self.section.id: (self.section, {}),
                            self.block.id: (self.block, {
                                self.center.id: (self.center, {}),
                            }),
                        },
                    ),
                    self.county.id: (
                        self.county,
                        {self.city.id: (self.city, {})},
                    ),
                },
            ),
        }
        self.assertEqual(hierarchy, desired_hierarchy)

@@ -171,16 +171,34 @@
 stCase):
+%0A%0A    @classmethod
 %0A    def
@@ -203,21 +203,25 @@
 ef setUp
-(self
+Class(cls
 ):%0A     
@@ -219,28 +219,27 @@
 s):%0A        
-self
+cls
 .domain = cr
@@ -277,253 +277,202 @@
     
-    def make_loc_type(name, parent_type=None):%0A            return LocationType.objects.create(%0A                domain=self.domain.name,%0A                name=name,%0A                code=name,%0A                parent_type=parent_type,%0A            )%0A
+@classmethod%0A    def tearDownClass(cls):%0A        cls.domain.delete()%0A%0A    def tearDown(self):%0A        LocationType.objects.filter(domain=self.domain.name).delete()%0A%0A    def test_hierarchy(self):
 %0A   
@@ -480,21 +480,16 @@
     
-self.
 state = 
 make
@@ -484,16 +484,21 @@
 state = 
+self.
 make_loc
@@ -513,37 +513,32 @@
 tate')%0A%0A        
-self.
 district = make_
@@ -532,16 +532,21 @@
 trict = 
+self.
 make_loc
@@ -563,21 +563,16 @@
 trict', 
-self.
 state)%0A 
@@ -574,29 +574,24 @@
 te)%0A        
-self.
 section = ma
@@ -588,16 +588,21 @@
 ction = 
+self.
 make_loc
@@ -618,21 +618,16 @@
 ction', 
-self.
 district
@@ -640,21 +640,16 @@
     
-self.
 block = 
 make
@@ -644,16 +644,21 @@
 block = 
+self.
 make_loc
@@ -672,21 +672,16 @@
 block', 
-self.
 district
@@ -686,29 +686,24 @@
 ct)%0A        
-self.
 center = mak
@@ -699,16 +699,21 @@
 enter = 
+self.
 make_loc
@@ -728,21 +728,16 @@
 enter', 
-self.
 block)%0A%0A
@@ -740,29 +740,24 @@
 k)%0A%0A        
-self.
 county = mak
@@ -749,24 +749,29 @@
    county = 
+self.
 make_loc_typ
@@ -782,21 +782,16 @@
 ounty', 
-self.
 state)%0A 
@@ -801,20 +801,20 @@
     
-self.
 city = 
+self.
 make
@@ -831,21 +831,16 @@
 'city', 
-self.
 county)%0A
@@ -844,162 +844,8 @@
 y)%0A%0A
-    def tearDown(self):%0A        self.domain.delete()%0A        LocationType.objects.filter(domain=self.domain.name).delete()%0A%0A    def test_hierarchy(self):%0A
     
@@ -952,29 +952,24 @@
             
-self.
 state.id: (%0A
@@ -984,21 +984,16 @@
         
-self.
 state,%0A 
@@ -1021,37 +1021,32 @@
                 
-self.
 district.id: (%0A 
@@ -1068,21 +1068,16 @@
         
-self.
 district
@@ -1132,21 +1132,16 @@
         
-self.
 section.
@@ -1145,21 +1145,16 @@
 on.id: (
-self.
 section,
@@ -1187,21 +1187,16 @@
         
-self.
 block.id
@@ -1198,21 +1198,16 @@
 ck.id: (
-self.
 block, %7B
@@ -1239,21 +1239,16 @@
         
-self.
 center.i
@@ -1251,21 +1251,16 @@
 er.id: (
-self.
 center, 
@@ -1362,29 +1362,24 @@
             
-self.
 county.id: (
@@ -1403,21 +1403,16 @@
         
-self.
 county,%0A
@@ -1436,21 +1436,16 @@
        %7B
-self.
 city.id:
@@ -1446,21 +1446,16 @@
 ty.id: (
-self.
 city, %7B%7D
@@ -1580,8 +1580,236 @@
 rarchy)%0A
+%0A    def make_loc_type(self, name, parent_type=None):%0A        return LocationType.objects.create(%0A            domain=self.domain.name,%0A            name=name,%0A            code=name,%0A            parent_type=parent_type,%0A        )%0A

clean up location types test