Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

Mock function to return different value based on argument passed

Tags:

python

mocking

In my test function, my mocks are not behaving properly because when I mock my get_config function, I only know how to mock it to return one value. How can I add some logic to my mock in order for the mock to return a different value only when get_config is passed the argument of "restricted"

def func():
    conf1 = get_config("conf1")
    conf2 = get_config("conf2")
    conf3 = get_config("conf3")
    restricted_datasets = get_config("restricted")
    dataset = get_config("dataset")
    if dataset not in restricted_datas:
        return run_code()
like image 950
Liondancer Avatar asked Jun 01 '26 05:06

Liondancer


2 Answers

You can assign a function to side_effect as described in official doc

from unittest.mock import Mock

def side_effect(value):
    return value

m = Mock(side_effect=side_effect)
m('restricted')
'restricted'
like image 160
Mauro Baraldi Avatar answered Jun 03 '26 18:06

Mauro Baraldi


class Ctx():
  def get_config(confName):
    return confName

def mock_get_config(value):
  if value == "conf1":
    return "confA"
  elif value == "conf2":
    return "confB"
  else:
    return "UnknownValue"

class CtxSourceFileTrial(unittest.TestCase):
  def test(self):
    mockCtx = Mock()
    mockCtx.get_config.side_effect = mock_get_config
    self.assertEqual("confA", mockCtx.get_config("conf1"))
    self.assertEqual("confB", mockCtx.get_config("conf2"))

#
# By the way I think Python is EXTREMELY screwy, Adligo
#
like image 28
user1303800 Avatar answered Jun 03 '26 17:06

user1303800