from os import getenv from unittest.mock import Mock, patch from core.models import Account, User # Create patch mixin to mock out the Elastic client class ElasticMock: """ Mixin to mock out the Elastic client. Patches the initialise_elasticsearch function to return a mock client. This client returns a mock response for the index function. """ @classmethod def setUpClass(cls): super(ElasticMock, cls).setUpClass() cls.patcher = patch("core.lib.elastic.initialise_elasticsearch") # Pretend the object has been created patcher = cls.patcher.start() fake_client = Mock() fake_client.index = Mock() fake_client.index.return_value = {"result": "created"} patcher.return_value = fake_client @classmethod def tearDownClass(cls): super(ElasticMock, cls).tearDownClass() cls.patcher.stop() class LiveBase: @classmethod def tearDownClass(cls): cls.account.client.close_all_positions() super(LiveBase, cls).tearDownClass() @classmethod def setUpClass(cls): super(LiveBase, cls).setUpClass() cls.live = getenv("LIVE", "0") == "1" account_name = getenv("TEST_ACCOUNT_NAME", "Test Account") exchange = getenv("TEST_ACCOUNT_EXCHANGE", None) api_key = getenv("TEST_ACCOUNT_API_KEY", None) api_secret = getenv("TEST_ACCOUNT_API_SECRET", None) if not cls.live: cls.skipTest(LiveBase, "Not running live tests") # Make sure all the variables are set cls.fail = False if not exchange: cls.fail = True reason = "No exchange specified" if not api_key: cls.fail = True reason = "No API key specified" if not api_secret: cls.fail = True reason = "No API secret specified" if cls.fail: print("Live tests aborted.") print( f""" Please check that the following environment variables are set: TEST_ACCOUNT_NAME="Test Account" TEST_ACCOUNT_EXCHANGE="oanda" TEST_ACCOUNT_API_KEY="xxx-xxx-xxxxxxxx-xxx" TEST_ACCOUNT_API_SECRET="xxx-xxx" If you have done this, please see the following line for more information: {reason} """ ) cls.user = User.objects.create_user( username="testuser", email="test@example.com", ) cls.account = Account.objects.create( user=cls.user, name=account_name, exchange=exchange, api_key=api_key, api_secret=api_secret, initial_balance=100, ) def setUp(self): if self.fail: self.skipTest("Live tests aborted")