Skip to content

Commit b9c2c4c

Browse files
iamavichal-geekiamavichal-geek
committed
[connection] add test for CommandTimeoutException in launch_command task
1 parent 55df52f commit b9c2c4c

File tree

1 file changed

+131
-108
lines changed

1 file changed

+131
-108
lines changed

‎openwisp_controller/connection/tests/test_tasks.py‎

Lines changed: 131 additions & 108 deletions
Original file line numberDiff line numberDiff line change
@@ -1,108 +1,131 @@
1-
import uuid
2-
from contextlib import redirect_stderr
3-
from io import StringIO
4-
from unittest import mock
5-
6-
from celery.exceptions import SoftTimeLimitExceeded
7-
from django.test import TestCase, TransactionTestCase
8-
from swapper import load_model
9-
10-
from ...config.tests.test_controller import TestRegistrationMixin
11-
from .. import tasks
12-
from .utils import CreateConnectionsMixin
13-
14-
Command = load_model("connection", "Command")
15-
OrganizationConfigSettings = load_model("config", "OrganizationConfigSettings")
16-
17-
18-
class TestTasks(CreateConnectionsMixin, TestCase):
19-
_mock_execute = "openwisp_controller.connection.base.models.AbstractCommand.execute"
20-
_mock_connect = (
21-
"openwisp_controller.connection.base.models.AbstractDeviceConnection.connect"
22-
)
23-
24-
@mock.patch("logging.Logger.warning")
25-
@mock.patch("time.sleep")
26-
def test_update_config_missing_config(self, mocked_sleep, mocked_warning):
27-
pk = self._create_device().pk
28-
tasks.update_config.delay(pk)
29-
mocked_warning.assert_called_with(
30-
f'update_config("{pk}") failed: Device has no config.'
31-
)
32-
mocked_sleep.assert_called_once()
33-
34-
@mock.patch("logging.Logger.warning")
35-
@mock.patch("time.sleep")
36-
def test_update_config_missing_device(self, mocked_sleep, mocked_warning):
37-
pk = uuid.uuid4()
38-
tasks.update_config.delay(pk)
39-
mocked_warning.assert_called_with(
40-
f'update_config("{pk}") failed: Device matching query does not exist.'
41-
)
42-
mocked_sleep.assert_called_once()
43-
44-
@mock.patch("logging.Logger.warning")
45-
def test_launch_command_missing(self, mocked_warning):
46-
pk = uuid.uuid4()
47-
tasks.launch_command.delay(pk)
48-
mocked_warning.assert_called_with(
49-
f'launch_command("{pk}") failed: Command matching query does not exist.'
50-
)
51-
52-
@mock.patch(_mock_execute, side_effect=SoftTimeLimitExceeded())
53-
@mock.patch(_mock_connect, return_value=True)
54-
def test_launch_command_timeout(self, *args):
55-
dc = self._create_device_connection()
56-
command = Command(
57-
device=dc.device,
58-
connection=dc,
59-
type="custom",
60-
input={"command": "/usr/sbin/exotic_command"},
61-
)
62-
command.full_clean()
63-
command.save()
64-
# must call this explicitly because lack of transactions in this test case
65-
tasks.launch_command.delay(command.pk)
66-
command.refresh_from_db()
67-
self.assertEqual(command.status, "failed")
68-
self.assertEqual(command.output, "Background task time limit exceeded.\n")
69-
70-
@mock.patch(_mock_execute, side_effect=RuntimeError("test error"))
71-
@mock.patch(_mock_connect, return_value=True)
72-
def test_launch_command_exception(self, *args):
73-
dc = self._create_device_connection()
74-
command = Command(
75-
device=dc.device,
76-
connection=dc,
77-
type="custom",
78-
input={"command": "/usr/sbin/exotic_command"},
79-
)
80-
command.full_clean()
81-
command.save()
82-
# must call this explicitly because lack of transactions in this test case
83-
with redirect_stderr(StringIO()) as stderr:
84-
tasks.launch_command.delay(command.pk)
85-
expected = f"An exception was raised while executing command {command.pk}"
86-
self.assertIn(expected, stderr.getvalue())
87-
command.refresh_from_db()
88-
self.assertEqual(command.status, "failed")
89-
self.assertEqual(command.output, "Internal system error: test error\n")
90-
91-
92-
class TestTransactionTasks(
93-
TestRegistrationMixin, CreateConnectionsMixin, TransactionTestCase
94-
):
95-
@mock.patch.object(tasks.update_config, "delay")
96-
def test_update_config_hostname_changed_on_reregister(self, mocked_update_config):
97-
device = self._create_device_config()
98-
self._create_device_connection(device=device)
99-
# Trigger re-registration with new hostname
100-
response = self.client.post(
101-
self.register_url,
102-
self._get_reregistration_payload(
103-
device,
104-
name="new-hostname",
105-
),
106-
)
107-
self.assertEqual(response.status_code, 201)
108-
mocked_update_config.assert_not_called()
1+
import uuid
2+
from contextlib import redirect_stderr
3+
from io import StringIO
4+
from unittest import mock
5+
6+
from celery.exceptions import SoftTimeLimitExceeded
7+
from django.test import TestCase, TransactionTestCase
8+
from swapper import load_model
9+
10+
from ...config.tests.test_controller import TestRegistrationMixin
11+
from .. import tasks
12+
from ..connectors.exceptions import CommandTimeoutException
13+
from .utils import CreateConnectionsMixin
14+
15+
Command = load_model("connection", "Command")
16+
OrganizationConfigSettings = load_model("config", "OrganizationConfigSettings")
17+
18+
19+
class TestTasks(CreateConnectionsMixin, TestCase):
20+
_mock_execute = "openwisp_controller.connection.base.models.AbstractCommand.execute"
21+
_mock_connect = (
22+
"openwisp_controller.connection.base.models.AbstractDeviceConnection.connect"
23+
)
24+
25+
@mock.patch("logging.Logger.warning")
26+
@mock.patch("time.sleep")
27+
def test_update_config_missing_config(self, mocked_sleep, mocked_warning):
28+
pk = self._create_device().pk
29+
tasks.update_config.delay(pk)
30+
mocked_warning.assert_called_with(
31+
f'update_config("{pk}") failed: Device has no config.'
32+
)
33+
mocked_sleep.assert_called_once()
34+
35+
@mock.patch("logging.Logger.warning")
36+
@mock.patch("time.sleep")
37+
def test_update_config_missing_device(self, mocked_sleep, mocked_warning):
38+
pk = uuid.uuid4()
39+
tasks.update_config.delay(pk)
40+
mocked_warning.assert_called_with(
41+
f'update_config("{pk}") failed: Device matching query does not exist.'
42+
)
43+
mocked_sleep.assert_called_once()
44+
45+
@mock.patch("logging.Logger.warning")
46+
def test_launch_command_missing(self, mocked_warning):
47+
pk = uuid.uuid4()
48+
tasks.launch_command.delay(pk)
49+
mocked_warning.assert_called_with(
50+
f'launch_command("{pk}") failed: Command matching query does not exist.'
51+
)
52+
53+
@mock.patch(_mock_execute, side_effect=SoftTimeLimitExceeded())
54+
@mock.patch(_mock_connect, return_value=True)
55+
def test_launch_command_timeout(self, *args):
56+
dc = self._create_device_connection()
57+
command = Command(
58+
device=dc.device,
59+
connection=dc,
60+
type="custom",
61+
input={"command": "/usr/sbin/exotic_command"},
62+
)
63+
command.full_clean()
64+
command.save()
65+
# must call this explicitly because lack of transactions in this test case
66+
tasks.launch_command.delay(command.pk)
67+
command.refresh_from_db()
68+
self.assertEqual(command.status, "failed")
69+
self.assertEqual(command.output, "Background task time limit exceeded.\n")
70+
71+
@mock.patch(
72+
_mock_execute,
73+
side_effect=CommandTimeoutException("connection timed out after 30s"),
74+
)
75+
@mock.patch(_mock_connect, return_value=True)
76+
def test_launch_command_ssh_timeout(self, *args):
77+
dc = self._create_device_connection()
78+
command = Command(
79+
device=dc.device,
80+
connection=dc,
81+
type="custom",
82+
input={"command": "/usr/sbin/exotic_command"},
83+
)
84+
command.full_clean()
85+
command.save()
86+
# must call this explicitly because lack of transactions in this test case
87+
tasks.launch_command.delay(command.pk)
88+
command.refresh_from_db()
89+
self.assertEqual(command.status, "failed")
90+
self.assertIn("The command took longer than expected", command.output)
91+
self.assertIn("connection timed out after 30s", command.output)
92+
93+
@mock.patch(_mock_execute, side_effect=RuntimeError("test error"))
94+
@mock.patch(_mock_connect, return_value=True)
95+
def test_launch_command_exception(self, *args):
96+
dc = self._create_device_connection()
97+
command = Command(
98+
device=dc.device,
99+
connection=dc,
100+
type="custom",
101+
input={"command": "/usr/sbin/exotic_command"},
102+
)
103+
command.full_clean()
104+
command.save()
105+
# must call this explicitly because lack of transactions in this test case
106+
with redirect_stderr(StringIO()) as stderr:
107+
tasks.launch_command.delay(command.pk)
108+
expected = f"An exception was raised while executing command {command.pk}"
109+
self.assertIn(expected, stderr.getvalue())
110+
command.refresh_from_db()
111+
self.assertEqual(command.status, "failed")
112+
self.assertEqual(command.output, "Internal system error: test error\n")
113+
114+
115+
class TestTransactionTasks(
116+
TestRegistrationMixin, CreateConnectionsMixin, TransactionTestCase
117+
):
118+
@mock.patch.object(tasks.update_config, "delay")
119+
def test_update_config_hostname_changed_on_reregister(self, mocked_update_config):
120+
device = self._create_device_config()
121+
self._create_device_connection(device=device)
122+
# Trigger re-registration with new hostname
123+
response = self.client.post(
124+
self.register_url,
125+
self._get_reregistration_payload(
126+
device,
127+
name="new-hostname",
128+
),
129+
)
130+
self.assertEqual(response.status_code, 201)
131+
mocked_update_config.assert_not_called()

0 commit comments

Comments
 (0)