mypointvpn/src/App/tests_services.py

from unittest import mock

import serverus
from django.test import TestCase
from Telegram.models import User as TelegramUser

from . import services
from .models import Point, Server


class FakeRunner:
	"""Подменяет запуск плейбуков serverus. fail - плейбук, который завершится ошибкой."""

	def __init__(self, fail: str | None = None):
		self.calls = []
		self.fail = fail

	def __call__(self, playbook, host, extravars):
		self.calls.append((playbook, host, extravars))
		if playbook == self.fail:
			return serverus.Result(ok=False, status='failed', error='boom', known_hosts=f'kh {host.address}')
		data = {'facts': {'uname': 'Linux test'}, 'xui': {'api_token': 'token'}}.get(playbook, {})
		return serverus.Result(ok=True, status='successful', data=data, known_hosts=f'kh {host.address}')

	def playbooks(self, address):
		return [playbook for playbook, host, _ in self.calls if host.address == address]

	def call(self, playbook, address):
		return next((host, extravars) for name, host, extravars in self.calls
			if name == playbook and host.address == address)


def patch_runner(runner):
	return mock.patch('serverus.runner.run_playbook', runner)


def create_point(owner=None, **kwargs):
	owner = owner or TelegramUser.objects.create(tg_id=1)
	return services.create_point(
		owner, services.Credentials('10.0.0.1', 'panel-pw'), services.Credentials('10.0.0.2', 'mask-pw'), **kwargs,
	)


class CreatePointTests(TestCase):
	def test_creates_point_with_servers_and_panel_access(self):
		point = create_point(domain='example.com')
		self.assertEqual(point.status, Point.Status.INSTALLING)
		self.assertEqual((point.panel.address, point.panel.root_password), ('10.0.0.1', 'panel-pw'))
		self.assertEqual((point.mask.address, point.mask.root_password), ('10.0.0.2', 'mask-pw'))
		low, high = services.PANEL_PORT_RANGE
		self.assertTrue(low <= point.panel_port < high)
		self.assertTrue(point.panel_username and point.panel_password and point.panel_path)


class InstallTests(TestCase):
	def test_install_point(self):
		point = create_point(domain='example.com')
		runner = FakeRunner()
		with patch_runner(runner):
			services.install_point(point)

		point.refresh_from_db()
		self.assertEqual(point.status, Point.Status.READY)
		self.assertEqual(point.panel_api_token, 'token')
		self.assertEqual(runner.playbooks('10.0.0.1'), ['bootstrap', 'xui', 'facts'])
		self.assertEqual(runner.playbooks('10.0.0.2'), ['bootstrap', 'proxy', 'facts'])

		host, extravars = runner.call('bootstrap', '10.0.0.1')
		self.assertEqual((host.port, host.known_hosts), (22, ''))
		self.assertEqual(extravars['login_password'], 'panel-pw')
		panel = point.panel
		self.assertEqual(extravars['new_password'], panel.root_password)
		self.assertEqual(extravars['ssh_port'], panel.ssh_port)
		self.assertEqual(extravars['public_key'], panel.public_key)
		self.assertEqual(panel.known_hosts, 'kh 10.0.0.1')
		self.assertEqual(panel.uname, 'Linux test')

		host, extravars = runner.call('xui', '10.0.0.1')
		self.assertEqual((host.port, host.known_hosts), (panel.ssh_port, 'kh 10.0.0.1'))
		self.assertEqual(extravars['domain'], 'example.com')
		self.assertEqual(extravars['panel_port'], point.panel_port)

		_, extravars = runner.call('proxy', '10.0.0.2')
		self.assertEqual(extravars['target_ip'], '10.0.0.1')

	def test_install_failure_marks_point(self):
		point = create_point()
		runner = FakeRunner(fail='xui')
		with patch_runner(runner), self.assertRaises(services.ProvisionError):
			services.install_point(point)
		point.refresh_from_db()
		self.assertEqual(point.status, Point.Status.ERROR)
		self.assertIn('boom', point.error)
		# после успешного bootstrap сервер уже на новом порту с новым паролем
		self.assertNotEqual(point.panel.ssh_port, 22)
		self.assertNotEqual(point.panel.root_password, 'panel-pw')

	def test_bootstrap_failure_keeps_password_and_key(self):
		point = create_point()
		with patch_runner(FakeRunner(fail='bootstrap')), self.assertRaises(services.ProvisionError):
			services.install_point(point)
		panel = point.panel
		self.assertEqual((panel.root_password, panel.ssh_port), ('panel-pw', 22))
		private_key = panel.private_key
		self.assertTrue(private_key)
		with patch_runner(FakeRunner()):
			services.install_point(point)
		self.assertEqual(Server.objects.get(pk=panel.pk).private_key, private_key)


class ReinstallTests(TestCase):
	def setUp(self):
		self.point = create_point()
		with patch_runner(FakeRunner()):
			services.install_point(self.point)

	def test_reinstall_one_server_with_new_password(self):
		runner = FakeRunner()
		with patch_runner(runner):
			services.reinstall(self.point, '10.0.0.2', 'new-provider-pw')
		self.assertEqual(runner.playbooks('10.0.0.1'), [])
		host, extravars = runner.call('bootstrap', '10.0.0.2')
		self.assertEqual((host.port, host.known_hosts), (22, ''))
		self.assertEqual(extravars['login_password'], 'new-provider-pw')
		self.point.refresh_from_db()
		self.assertEqual(self.point.status, Point.Status.READY)

	def test_reinstall_all_with_old_passwords(self):
		old_password = self.point.panel.root_password
		runner = FakeRunner()
		with patch_runner(runner):
			services.reinstall(self.point)
		_, extravars = runner.call('bootstrap', '10.0.0.1')
		self.assertEqual(extravars['login_password'], old_password)
		self.assertTrue(runner.playbooks('10.0.0.2'))


class ReleaseTests(TestCase):
	def setUp(self):
		self.point = create_point()
		with patch_runner(FakeRunner()):
			services.install_point(self.point)

	def test_release_deletes_point_and_returns_passwords(self):
		runner = FakeRunner()
		with patch_runner(runner):
			passwords, errors = services.release_point(self.point)
		self.assertEqual(errors, {})
		self.assertEqual(passwords['10.0.0.1'], runner.call('release', '10.0.0.1')[1]['new_password'])
		self.assertEqual(set(passwords), {'10.0.0.1', '10.0.0.2'})
		self.assertFalse(Point.objects.exists())
		self.assertFalse(Server.objects.exists())

	def test_partial_release_keeps_point(self):
		runner = FakeRunner(fail='release')
		with patch_runner(runner):
			passwords, errors = services.release_point(self.point)
		self.assertEqual(passwords, {})
		self.assertEqual(set(errors), {'10.0.0.1', '10.0.0.2'})
		self.point.refresh_from_db()
		self.assertEqual(self.point.status, Point.Status.ERROR)