2023-12-11 02:20:34 -05:00
"""
A test script to test the gRPC service
"""
import unittest
import subprocess
import time
import backend_pb2
import backend_pb2_grpc
import grpc
class TestBackendServicer ( unittest . TestCase ) :
"""
TestBackendServicer is the class that tests the gRPC service
"""
def setUp ( self ) :
"""
This method sets up the gRPC service by starting the server
"""
2024-05-11 11:32:46 -05:00
self . service = subprocess . Popen ( [ " python3 " , " backend.py " , " --addr " , " localhost:50051 " ] )
2023-12-11 02:20:34 -05:00
def tearDown ( self ) - > None :
"""
This method tears down the gRPC service by terminating the server
"""
self . service . kill ( )
self . service . wait ( )
def test_server_startup ( self ) :
"""
This method tests if the server starts up successfully
"""
time . sleep ( 10 )
try :
self . setUp ( )
with grpc . insecure_channel ( " localhost:50051 " ) as channel :
stub = backend_pb2_grpc . BackendStub ( channel )
response = stub . Health ( backend_pb2 . HealthMessage ( ) )
self . assertEqual ( response . message , b ' OK ' )
except Exception as err :
print ( err )
self . fail ( " Server failed to start " )
finally :
self . tearDown ( )
def test_load_model ( self ) :
"""
This method tests if the model is loaded successfully
"""
time . sleep ( 10 )
try :
self . setUp ( )
with grpc . insecure_channel ( " localhost:50051 " ) as channel :
stub = backend_pb2_grpc . BackendStub ( channel )
2024-08-30 10:04:36 +02:00
response = stub . LoadModel ( backend_pb2 . ModelOptions ( Model = " Lykon/dreamshaper-8 " ) )
2023-12-11 02:20:34 -05:00
self . assertTrue ( response . success )
self . assertEqual ( response . message , " Model loaded successfully " )
except Exception as err :
print ( err )
self . fail ( " LoadModel service failed " )
finally :
self . tearDown ( )
def test ( self ) :
"""
This method tests if the backend can generate images
"""
time . sleep ( 10 )
try :
self . setUp ( )
with grpc . insecure_channel ( " localhost:50051 " ) as channel :
stub = backend_pb2_grpc . BackendStub ( channel )
2024-08-30 10:04:36 +02:00
response = stub . LoadModel ( backend_pb2 . ModelOptions ( Model = " Lykon/dreamshaper-8 " ) )
2023-12-11 02:20:34 -05:00
print ( response . message )
self . assertTrue ( response . success )
image_req = backend_pb2 . GenerateImageRequest ( positive_prompt = " cat " , width = 16 , height = 16 , dst = " test.jpg " )
re = stub . GenerateImage ( image_req )
self . assertTrue ( re . success )
except Exception as err :
print ( err )
self . fail ( " Image gen service failed " )
finally :
2024-08-30 10:04:36 +02:00
self . tearDown ( )
