more tests

Author: yevgenypats

cloudquery/sdk/internal/servers/plugin_v3/plugin.py

@@ -16,7 +16,7 @@ def GetName(self, request, context):
         return plugin_pb2.GetName.Response(name=self._plugin.name())
 
     def GetVersion(self, request, context):
-        return plugin_pb2.GetVersion.Response(name=self._plugin.version())
+        return plugin_pb2.GetVersion.Response(version=self._plugin.version())
 
     def Init(self, request: plugin_pb2.Init.Request, context):
         self._plugin.init(request.spec)

cloudquery/sdk/serve/plugin.py

@@ -63,15 +63,18 @@ def run(self, args):
 
     def _serve(self, args):
         logger = get_logger(args)
-        server = grpc.server(futures.ThreadPoolExecutor(max_workers=10))
+        self._server = grpc.server(futures.ThreadPoolExecutor(max_workers=10))
         discovery_pb2_grpc.add_DiscoveryServicer_to_server(
-            DiscoveryServicer([3]), server)
+            DiscoveryServicer([3]), self._server)
         plugin_pb2_grpc.add_PluginServicer_to_server(
-            PluginServicer(self._plugin, logger), server)
-        server.add_insecure_port(args.address)
+            PluginServicer(self._plugin, logger), self._server)
+        self._server.add_insecure_port(args.address)
         print("Starting server. Listening on " + args.address)
-        server.start()
-        server.wait_for_termination()
+        self._server.start()
+        self._server.wait_for_termination()
+    
+    def stop(self):
+        self._server.stop(5)
 
     def _generate_docs(self, args):
         print("Generating docs in format: " + args.format)

tests/serve/plugin.py

@@ -1,9 +1,28 @@
 
+import random
+import grpc
+import time
+from concurrent import futures
 from cloudquery.sdk import serve
 from cloudquery.sdk import plugin
+from cloudquery.plugin_v3 import plugin_pb2_grpc, plugin_pb2
+
 
 def test_plugin_serve():
-    pass
-    # p = plugin.Plugin("test", "v1.0.0")
-    # cmd = serve.PluginCommand(p)
-    # cmd.run(["plugin", "serve"])
+    p = plugin.Plugin("test", "v1.0.0")
+    cmd = serve.PluginCommand(p)
+    port = random.randint(5000, 50000)
+    pool = futures.ThreadPoolExecutor(max_workers=1)
+    pool.submit(cmd.run, ["serve", "--address", f"[::]:{port}"])
+    time.sleep(1)
+    try:
+      with grpc.insecure_channel(f'localhost:{port}') as channel:
+          stub = plugin_pb2_grpc.PluginStub(channel)
+          response = stub.GetName(plugin_pb2.GetName.Request())
+          assert response.name == "test"
+          
+          response = stub.GetVersion(plugin_pb2.GetVersion.Request())
+          assert response.version == "v1.0.0"
+    finally:
+      cmd.stop()
+      pool.shutdown()