Untitled

from typing import Optional, Dict, Any
from unittest.mock import patch
import logger

BEOWULF = "Beowulf"

# SOME CODE #
class InflationRiskViewCollectionBase:
    def __init__(self, buckets, desk):
        self.buckets = buckets
        self.desk = desk
        self.risk_sources = {}

    def load_data_from_beowulf(self, data_sources: Optional[Dict[str, Any]] = None) -> str:
        if data_sources is None:
            data_sources = dict()
        
        # Little explanation about the following two conditions
        # First of all, 'data_sources' is an override of what could be retrieved by 
        # BeowulfDataRetriever().fetch_data() function
        # BW = Beowulf
        # BW IN self.risk_sources and BW IN data_sources: override (by data_sources[BW])
        # BW NOT IN self.risk_sources and BW IN data_sources: override (really just assign)
        # BW NOT IN self.risk_sources and BW NOT IN data_sources: fetching
        # BW IN self.risk_sources and BW not in data_sources:
        #   BW stale : Re-fetching of BW data
        #   BW not stale : nothing happpens, just status update

        if BEOWULF not in data_sources:
            BW_in_risk_srcs = BEOWULF in self.risk_sources
            if not BW_in_risk_srcs or self.risk_sources[BEOWULF].is_stale:
                logger.debug("Fetching beowulf data")
                if BW_in_risk_srcs:
                    bw_status = "Refetching BW data"
                else:
                    bw_status = "Fetch BW data"
                try:
                    data_sources[BEOWULF] = BeowulfDataRetriever(
                        desk="TEST"
                    ).fetch_data()
                except (Exception, BeowulfDataRetrieverExeption) as err:
                    raise
            else:
                bw_status = "Data not stale, nothing to reload"
        else:
            bw_status="Beowulf data overridden"
        self.risk_sources.update(data_sources)
        return bw_status


# -----------------------------------------------------------------------------
@patch_decorator_factory(beowulf_patch, beowulf_response_patch)
def test_load_data_from_beowulf_case_3(self, bw_mock, bw_response_mock):
    expected_status = ""
    inf_base = InflationRiskViewCollectionBase(buckets=[], desk="TEST")
    with patch.object(
        BeowulfDataRetriever,
        "fetch_data"
    ) as mock_fetch_data:
        bw_response_mock.return_value.is_stale = True
        inf_base.risk_sources[BEOWULF] = bw_response_mock()
        actual = inf_base.load_data_from_beowulf()
    self.assertEqual(actual, expected_status) # works
    mock_fetch_data.assert_called_once() # DOES NOT WORK. CALLED 0 TIMES