diff --git a/utils/tests.py b/utils/tests.py index c4608e73..362a84bc 100644 --- a/utils/tests.py +++ b/utils/tests.py @@ -1,16 +1,19 @@ import uuid +from time import sleep from unittest.mock import patch import stripe +from celery.result import AsyncResult +from django.conf import settings from django.http.request import HttpRequest from django.test import Client -from django.test import TestCase +from django.test import TestCase, override_settings from model_mommy import mommy from datacenterlight.models import StripePlan from membership.models import StripeCustomer from utils.stripe_utils import StripeUtils -from django.conf import settings +from .tasks import save_ssh_key class BaseTestCase(TestCase): @@ -235,3 +238,29 @@ class StripePlanTestCase(TestStripeCustomerDescription): 'response_object').stripe_plan_id}]) self.assertIsNone(result.get('response_object'), None) self.assertIsNotNone(result.get('error')) + + +class SaveSSHKeyTestCase(TestCase): + """ + A test case to test the celery save_ssh_key task + """ + + @override_settings( + task_eager_propagates=True, + task_always_eager=True, + ) + def setUp(self): + self.public_key = ["This is a test", ] + self.hosts = ['localhost'] + + def test_save_ssh_key(self): + async_task = save_ssh_key.delay(self.hosts, self.public_key) + save_ssh_key_result = None + for i in range(0, 10): + sleep(5) + res = AsyncResult(async_task.task_id) + if type(res.result) is bool: + save_ssh_key_result = res.result + break + self.assertIsNotNone(save_ssh_key, "save_ssh_key_result is None") + self.assertTrue(save_ssh_key_result, "save_ssh_key_result is False")