| import unittest | |
| import subprocess | |
| import time | |
| import backend_pb2 | |
| import backend_pb2_grpc | |
| import grpc | |
| import unittest | |
| import subprocess | |
| import time | |
| import grpc | |
| import backend_pb2_grpc | |
| import backend_pb2 | |
| class TestBackendServicer(unittest.TestCase): | |
| """ | |
| TestBackendServicer is the class that tests the gRPC service. | |
| This class contains methods to test the startup and shutdown of the gRPC service. | |
| """ | |
| def setUp(self): | |
| self.service = subprocess.Popen(["python", "backend.py", "--addr", "localhost:50051"]) | |
| time.sleep(10) | |
| def tearDown(self) -> None: | |
| self.service.terminate() | |
| self.service.wait() | |
| def test_server_startup(self): | |
| 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 | |
| """ | |
| try: | |
| self.setUp() | |
| with grpc.insecure_channel("localhost:50051") as channel: | |
| stub = backend_pb2_grpc.BackendStub(channel) | |
| response = stub.LoadModel(backend_pb2.ModelOptions(Model="bigscience/bloom-560m")) | |
| print(response) | |
| 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() | |