168 lines
6.5 KiB
Python
168 lines
6.5 KiB
Python
import pytest
|
|
import uuid
|
|
import asyncio
|
|
from channels.testing import WebsocketCommunicator
|
|
from project.config.asgi import application
|
|
from project.remotectl.models import RemoteHost, BatchScript, CommandTask
|
|
|
|
class DummyConn:
|
|
def close(self):
|
|
pass
|
|
|
|
@pytest.mark.django_db(transaction=True)
|
|
@pytest.mark.asyncio
|
|
async def test_command_failure(monkeypatch, django_user_model):
|
|
user = await django_user_model.objects.acreate(username='userf')
|
|
host = await RemoteHost.objects.acreate(name='hf', hostname='localhost', username='u')
|
|
|
|
from project.remotectl import consumers
|
|
|
|
async def fake_open_connection(rh):
|
|
return DummyConn()
|
|
|
|
async def fake_run_command(conn, command, on_chunk, cancel_event):
|
|
await on_chunk('stdout', 'doing something\n')
|
|
return 5 # non-zero
|
|
|
|
monkeypatch.setattr(consumers, 'open_connection', fake_open_connection)
|
|
monkeypatch.setattr(consumers, 'run_command', fake_run_command)
|
|
|
|
session_id = uuid.uuid4()
|
|
comm = WebsocketCommunicator(application, f"/ws/ssh/{session_id}/stream/")
|
|
comm.scope['user'] = user
|
|
connected, _ = await comm.connect(); assert connected
|
|
await comm.receive_json_from() # connected
|
|
await comm.send_json_to({'action':'start','host_id':host.id,'command':'failingcmd'})
|
|
msg = await comm.receive_json_from(); assert msg['event']=='started'
|
|
# drain until completed
|
|
status = None
|
|
while True:
|
|
m = await comm.receive_json_from()
|
|
if m['event'] == 'completed':
|
|
status = m['status']; break
|
|
assert status == 'failed'
|
|
await comm.disconnect()
|
|
|
|
@pytest.mark.django_db(transaction=True)
|
|
@pytest.mark.asyncio
|
|
async def test_batch_success(monkeypatch, django_user_model):
|
|
user = await django_user_model.objects.acreate(username='userb')
|
|
host = await RemoteHost.objects.acreate(name='hb', hostname='localhost', username='u')
|
|
batch = await BatchScript.objects.acreate(name='b1', description='d', script='cd /tmp\necho one\necho two\n# comment\necho three')
|
|
|
|
from project.remotectl import consumers
|
|
async def fake_open_connection(rh):
|
|
return DummyConn()
|
|
|
|
calls = []
|
|
async def fake_run_command(conn, command, on_chunk, cancel_event):
|
|
calls.append(command)
|
|
await on_chunk('stdout', 'out\n')
|
|
return 0
|
|
|
|
monkeypatch.setattr(consumers, 'open_connection', fake_open_connection)
|
|
monkeypatch.setattr(consumers, 'run_command', fake_run_command)
|
|
|
|
session_id = uuid.uuid4()
|
|
comm = WebsocketCommunicator(application, f"/ws/ssh/{session_id}/stream/")
|
|
comm.scope['user'] = user
|
|
connected, _ = await comm.connect(); assert connected
|
|
await comm.receive_json_from() # connected
|
|
await comm.send_json_to({'action':'start_batch','host_id':host.id,'batch_id':batch.id})
|
|
msg = await comm.receive_json_from(); assert msg['event']=='started'
|
|
step_headers = 0
|
|
completed = None
|
|
while True:
|
|
m = await comm.receive_json_from()
|
|
if m['event']=='chunk' and m['data'].startswith('\n>>>'):
|
|
step_headers += 1
|
|
if m['event']=='completed':
|
|
completed = m; break
|
|
assert completed['status']=='ok'
|
|
# Steps: cd, echo one, echo two, echo three => 4 headers
|
|
assert step_headers == 5 or step_headers >= 4 # tolerate extra formatting
|
|
# run_command should have been called for 3 executable echo commands
|
|
assert len(calls) == 3
|
|
await comm.disconnect()
|
|
|
|
@pytest.mark.django_db(transaction=True)
|
|
@pytest.mark.asyncio
|
|
async def test_batch_failure(monkeypatch, django_user_model):
|
|
user = await django_user_model.objects.acreate(username='userbf')
|
|
host = await RemoteHost.objects.acreate(name='hbf', hostname='localhost', username='u')
|
|
batch = await BatchScript.objects.acreate(name='bf', description='d', script='echo ok\necho fail\necho skip')
|
|
|
|
from project.remotectl import consumers
|
|
async def fake_open_connection(rh):
|
|
return DummyConn()
|
|
|
|
call_count = 0
|
|
async def fake_run_command(conn, command, on_chunk, cancel_event):
|
|
nonlocal call_count
|
|
call_count += 1
|
|
await on_chunk('stdout', f'run {call_count}\n')
|
|
if call_count == 2:
|
|
return 9 # fail second step
|
|
return 0
|
|
|
|
monkeypatch.setattr(consumers, 'open_connection', fake_open_connection)
|
|
monkeypatch.setattr(consumers, 'run_command', fake_run_command)
|
|
|
|
session_id = uuid.uuid4()
|
|
comm = WebsocketCommunicator(application, f"/ws/ssh/{session_id}/stream/")
|
|
comm.scope['user'] = user
|
|
connected, _ = await comm.connect(); assert connected
|
|
await comm.receive_json_from() # connected
|
|
await comm.send_json_to({'action':'start_batch','host_id':host.id,'batch_id':batch.id})
|
|
await comm.receive_json_from() # started
|
|
while True:
|
|
m = await comm.receive_json_from()
|
|
if m['event']=='completed':
|
|
assert m['status'] == 'failed'
|
|
break
|
|
assert call_count == 2 # third not executed
|
|
await comm.disconnect()
|
|
|
|
@pytest.mark.django_db(transaction=True)
|
|
@pytest.mark.asyncio
|
|
async def test_cancel_command(monkeypatch, django_user_model):
|
|
user = await django_user_model.objects.acreate(username='userc')
|
|
host = await RemoteHost.objects.acreate(name='hc', hostname='localhost', username='u')
|
|
|
|
from project.remotectl import consumers
|
|
async def fake_open_connection(rh):
|
|
return DummyConn()
|
|
|
|
async def fake_run_command(conn, command, on_chunk, cancel_event):
|
|
# emit some output repeatedly until canceled
|
|
for i in range(5):
|
|
await on_chunk('stdout', f'line {i}\n')
|
|
await asyncio.sleep(0)
|
|
if cancel_event.is_set():
|
|
return 130
|
|
return 0
|
|
|
|
monkeypatch.setattr(consumers, 'open_connection', fake_open_connection)
|
|
monkeypatch.setattr(consumers, 'run_command', fake_run_command)
|
|
|
|
session_id = uuid.uuid4()
|
|
comm = WebsocketCommunicator(application, f"/ws/ssh/{session_id}/stream/")
|
|
comm.scope['user'] = user
|
|
connected, _ = await comm.connect(); assert connected
|
|
await comm.receive_json_from() # connected
|
|
await comm.send_json_to({'action':'start','host_id':host.id,'command':'longrun'})
|
|
await comm.receive_json_from() # started
|
|
# send cancel quickly
|
|
await comm.send_json_to({'action':'cancel'})
|
|
saw_canceling = False
|
|
status = None
|
|
while True:
|
|
m = await comm.receive_json_from()
|
|
if m['event']=='canceling':
|
|
saw_canceling = True
|
|
if m['event']=='completed':
|
|
status = m['status']; break
|
|
assert saw_canceling
|
|
assert status == 'canceled'
|
|
await comm.disconnect()
|