tests: add diffusers tests (#1419)

backend/python/diffusers/Makefile

@@ -9,3 +9,6 @@ run:
 	@echo "Running diffusers..."
 	bash run.sh
 	@echo "Diffusers run."
+
+test:
+	bash test.sh

backend/python/diffusers/backend_diffusers.py

@@ -303,12 +303,17 @@ class BackendServicer(backend_pb2_grpc.BackendServicer):
 
         prompt = request.positive_prompt
 
+        steps = 1
+
+        if request.step != 0:
+            steps = request.step
+
         # create a dictionary of values for the parameters
         options = {
             "negative_prompt":     request.negative_prompt, 
             "width":               request.width, 
             "height":              request.height,
-            "num_inference_steps": request.step,
+            "num_inference_steps": steps,
         }
 
         if request.src != "":

backend/python/diffusers/test.py

@@ -0,0 +1,84 @@
+"""
+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
+        """
+        self.service = subprocess.Popen(["python3", "backend_diffusers.py", "--addr", "localhost:50051"])
+
+    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)
+                response = stub.LoadModel(backend_pb2.ModelOptions(Model="runwayml/stable-diffusion-v1-5", PipelineType="StableDiffusionPipeline"))
+                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)
+                response = stub.LoadModel(backend_pb2.ModelOptions(Model="runwayml/stable-diffusion-v1-5", PipelineType="StableDiffusionPipeline"))
+                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:
+            self.tearDown()

backend/python/diffusers/test.sh

@@ -0,0 +1,14 @@
+#!/bin/bash
+
+##
+## A bash script wrapper that runs the diffusers server with conda
+
+export PATH=$PATH:/opt/conda/bin
+
+# Activate conda environment
+source activate diffusers
+
+# get the directory where the bash script is located
+DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )"
+
+python -m unittest $DIR/test.py

backend/python/sentencetransformers/test_sentencetransformers.py

@@ -24,7 +24,7 @@ class TestBackendServicer(unittest.TestCase):
         """
         This method tears down the gRPC service by terminating the server
         """
-        self.service.terminate()
+        self.service.kill()
         self.service.wait()
 
     def test_server_startup(self):

backend/python/transformers/test_transformers_server.py

@@ -24,7 +24,7 @@ class TestBackendServicer(unittest.TestCase):
         """
         This method tears down the gRPC service by terminating the server
         """
-        self.service.terminate()
+        self.service.kill()
         self.service.wait()
 
     def test_server_startup(self):