fisk/core/tests/helpers.py

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")