<s> from datetime import datetime <EOL> from corehq . apps . commtrack . models import StockState <EOL> from corehq . apps . domain . shortcuts import create_domain <EOL> from corehq . apps . reports . commtrack . data_sources import StockStatusDataSource <EOL> from corehq . apps . users . models import WebUser <EOL> from dimagi . utils . couch . loosechange import map_reduce <EOL> from corehq . apps . commtrack . helpers import make_product <EOL> from corehq . apps . commtrack . tests . util import make_loc <EOL> CURRENT_STOCK = StockStatusDataSource . SLUG_CURRENT_STOCK <EOL> PRODUCT_ID = StockStatusDataSource . SLUG_PRODUCT_ID <EOL> LOCATION_ID = StockStatusDataSource . SLUG_LOCATION_ID <EOL> TEST_DOMAIN = '<STR_LIT>' <EOL> class DataSourceTest ( object ) : <EOL> @ classmethod <EOL> def setUpClass ( cls ) : <EOL> cls . domain = create_domain ( TEST_DOMAIN ) <EOL> cls . couch_user = WebUser . create ( None , "<STR_LIT>" , "<STR_LIT>" ) <EOL> cls . couch_user . add_domain_membership ( TEST_DOMAIN , is_admin = True ) <EOL> cls . couch_user . save ( ) <EOL> cls . products = { <EOL> '<STR_LIT>' : make_product ( TEST_DOMAIN , '<STR_LIT>' , '<STR_LIT>' ) , <EOL> '<STR_LIT>' : make_product ( TEST_DOMAIN , '<STR_LIT>' , '<STR_LIT>' ) <EOL> } <EOL> test_setup = { <EOL> '<STR_LIT:A>' : { <EOL> '<STR_LIT>' : { <EOL> '<STR_LIT>' : { '<STR_LIT>' : <NUM_LIT:4> , '<STR_LIT>' : <NUM_LIT:0> } , <EOL> '<STR_LIT>' : { '<STR_LIT>' : <NUM_LIT:3> } , <EOL> } , <EOL> '<STR_LIT>' : { <EOL> '<STR_LIT>' : { '<STR_LIT>' : <NUM_LIT:2> } <EOL> } <EOL> } , <EOL> '<STR_LIT:B>' : { <EOL> '<STR_LIT>' : { <EOL> '<STR_LIT>' : { '<STR_LIT>' : <NUM_LIT:1> , '<STR_LIT>' : <NUM_LIT:1> } <EOL> } <EOL> } <EOL> } <EOL> cls . sites = { } <EOL> cls . regions = { } <EOL> cls . districts = { } <EOL> for region_name , districts in test_setup . items ( ) : <EOL> region = make_loc ( region_name , type = '<STR_LIT>' ) <EOL> cls . regions [ region_name ] = region <EOL> for district_name , sites in districts . items ( ) : <EOL> district = make_loc ( district_name , type = '<STR_LIT>' , parent = region ) <EOL> cls . districts [ district_name ] = district <EOL> for site_name , products in sites . items ( ) : <EOL> site = make_loc ( site_name , type = '<STR_LIT>' , parent = district , domain = TEST_DOMAIN ) <EOL> cls . sites [ site_name ] = ( site , products ) <EOL> supply_point = site . linked_supply_point ( ) <EOL> for p_code , stock in products . items ( ) : <EOL> prod = cls . products [ p_code ] <EOL> StockState . objects . create ( <EOL> section_id = '<STR_LIT>' , <EOL> case_id = supply_point . case_id , <EOL> product_id = prod . _id , <EOL> stock_on_hand = stock , <EOL> last_modified_date = datetime . utcnow ( ) , <EOL> ) <EOL> @ classmethod <EOL> def tearDownClass ( cls ) : <EOL> cls . couch_user . delete ( ) <EOL> cls . domain . delete ( ) <EOL> def test_raw_cases ( self ) : <EOL> config = { <EOL> '<STR_LIT>' : TEST_DOMAIN <EOL> } <EOL> data = list ( StockStatusDataSource ( config ) . get_data ( ) ) <EOL> self . assertEqual ( len ( data ) , <NUM_LIT:6> ) <EOL> by_location = map_reduce ( lambda row : [ ( row [ LOCATION_ID ] , ) ] , data = data , include_docs = True ) <EOL> for site , products in self . sites . values ( ) : <EOL> site_id = site . _id <EOL> rows = by_location [ site_id ] <EOL> by_product = dict ( ( row [ PRODUCT_ID ] , row ) for row in rows ) <EOL> for code , level in products . items ( ) : <EOL> product_id = self . products [ code ] . _id <EOL> self . assertEqual ( by_product [ product_id ] [ CURRENT_STOCK ] , level ) <EOL> def test_raw_cases_location ( self ) : <EOL> location = self . districts [ '<STR_LIT>' ] . _id <EOL> config = { <EOL> '<STR_LIT>' : TEST_DOMAIN , <EOL> '<STR_LIT>' : location <EOL> } <EOL> data = list ( StockStatusDataSource ( config ) . get_data ( ) ) <EOL> self . assertEqual ( len ( data ) , <NUM_LIT:1> ) <EOL> self . assertEqual ( data [ <NUM_LIT:0> ] [ LOCATION_ID ] , self . sites [ '<STR_LIT>' ] [ <NUM_LIT:0> ] . _id ) <EOL> self . assertEqual ( data [ <NUM_LIT:0> ] [ PRODUCT_ID ] , self . products [ '<STR_LIT>' ] . _id ) <EOL> self . assertEqual ( data [ <NUM_LIT:0> ] [ CURRENT_STOCK ] , <NUM_LIT:2> ) <EOL> def test_aggregate_level1 ( self ) : <EOL> location = self . regions [ '<STR_LIT:A>' ] . _id <EOL> config = { <EOL> '<STR_LIT>' : TEST_DOMAIN , <EOL> '<STR_LIT>' : location , <EOL> '<STR_LIT>' : True <EOL> } <EOL> data = list ( StockStatusDataSource ( config ) . get_data ( ) ) <EOL> self . assertEqual ( len ( data ) , <NUM_LIT:2> ) <EOL> by_product = dict ( ( row [ PRODUCT_ID ] , row ) for row in data ) <EOL> pA_id = self . products [ '<STR_LIT>' ] . _id <EOL> pB_id = self . products [ '<STR_LIT>' ] . _id <EOL> self . assertEqual ( by_product [ pA_id ] [ CURRENT_STOCK ] , <NUM_LIT:6> ) <EOL> self . assertEqual ( by_product [ pB_id ] [ CURRENT_STOCK ] , <NUM_LIT:3> ) <EOL> def test_aggregate_level2 ( self ) : <EOL> location = self . districts [ '<STR_LIT>' ] . _id <EOL> config = { <EOL> '<STR_LIT>' : TEST_DOMAIN , <EOL> '<STR_LIT>' : location , <EOL> '<STR_LIT>' : True <EOL> } <EOL> data = list ( StockStatusDataSource ( config ) . get_data ( ) ) <EOL> self . assertEqual ( len ( data ) , <NUM_LIT:1> ) <EOL> self . assertEqual ( data [ <NUM_LIT:0> ] [ CURRENT_STOCK ] , <NUM_LIT:2> ) </s>
