from django.test import TestCase
|
from ml_model_providers.models import ModelProviderConnection
|
from organizations.tests.factories import OrganizationFactory
|
|
|
class TestModelProviderConnection(TestCase):
|
@classmethod
|
def setUpTestData(cls):
|
cls.organization = OrganizationFactory()
|
|
def test_has_reached_budget_limit(self):
|
mpc = ModelProviderConnection.objects.create(
|
organization=self.organization,
|
created_by=self.organization.created_by,
|
)
|
# Not internal
|
assert not mpc.has_reached_budget_limit()
|
mpc.is_internal = True
|
mpc.save()
|
# No budget_total_spent
|
assert not mpc.has_reached_budget_limit()
|
mpc.budget_total_spent = 1000
|
mpc.save()
|
# No budget_limit
|
assert not mpc.has_reached_budget_limit()
|
mpc.budget_limit = 2000
|
mpc.save()
|
# Not reached
|
assert not mpc.has_reached_budget_limit()
|
mpc.budget_total_spent = 2001
|
# Reached
|
assert mpc.has_reached_budget_limit()
|
