skynet/tests/test_skynet.py

71 lines
1.8 KiB
Python
Raw Normal View History

#!/usr/bin/python
import logging
import trio
2022-12-11 16:06:07 +00:00
import pynng
import pytest
import trio_asyncio
from skynet.brain import run_skynet
from skynet.structs import *
from skynet.frontend import open_skynet_rpc
async def test_skynet(skynet_running):
...
2022-12-11 16:06:07 +00:00
async def test_skynet_attempt_insecure(skynet_running):
with pytest.raises(pynng.exceptions.NNGException) as e:
async with open_skynet_rpc('bad-actor'):
...
2022-12-11 16:06:07 +00:00
assert str(e.value) == 'Connection shutdown'
async def test_skynet_dgpu_connection_simple(skynet_running):
2022-12-11 16:06:07 +00:00
async with open_skynet_rpc(
'dgpu-0',
2022-12-11 16:06:07 +00:00
security=True,
cert_name='whitelist/testing',
key_name='testing'
) as rpc_call:
# check 0 nodes are connected
res = await rpc_call('dgpu_workers')
assert 'ok' in res.result
assert res.result['ok'] == 0
# check next worker is None
res = await rpc_call('dgpu_next')
assert 'ok' in res.result
assert res.result['ok'] == None
# connect 1 dgpu
res = await rpc_call('dgpu_online')
assert 'ok' in res.result
# check 1 node is connected
res = await rpc_call('dgpu_workers')
assert 'ok' in res.result
assert res.result['ok'] == 1
# check next worker is 0
res = await rpc_call('dgpu_next')
assert 'ok' in res.result
assert res.result['ok'] == 0
# disconnect 1 dgpu
res = await rpc_call('dgpu_offline')
assert 'ok' in res.result
# check 0 nodes are connected
res = await rpc_call('dgpu_workers')
assert 'ok' in res.result
assert res.result['ok'] == 0
# check next worker is None
res = await rpc_call('dgpu_next')
assert 'ok' in res.result
assert res.result['ok'] == None