diff --git a/src/opengeodeweb_viewer/rpc/schemas/take_screenshot.json b/src/opengeodeweb_viewer/rpc/schemas/take_screenshot.json index b663808..f292770 100644 --- a/src/opengeodeweb_viewer/rpc/schemas/take_screenshot.json +++ b/src/opengeodeweb_viewer/rpc/schemas/take_screenshot.json @@ -6,7 +6,11 @@ "type": "string" }, "output_extension": { - "type": "string" + "type": "string", + "enum": [ + "png", + "jpg" + ] }, "include_background": { "type": "boolean" diff --git a/src/opengeodeweb_viewer/vtk_protocol.py b/src/opengeodeweb_viewer/vtk_protocol.py index be1714a..bab744f 100644 --- a/src/opengeodeweb_viewer/vtk_protocol.py +++ b/src/opengeodeweb_viewer/vtk_protocol.py @@ -295,9 +295,8 @@ def takeScreenshot(self, params): w2if = vtkWindowToImageFilter() if not include_background: - # renderer.SetBackground([255,255,255]) - # renderer.SetLayer(1) renderWindow.SetAlphaBitPlanes(1) + w2if.SetInput(renderWindow) w2if.SetInputBufferTypeToRGBA() w2if.ReadFrontBufferOff() @@ -308,12 +307,19 @@ def takeScreenshot(self, params): writer = vtkPNGWriter() elif output_extension == "jpg": writer = vtkJPEGWriter() + else: + raise Exception("output_extension not supported") - writer.SetFileName(os.path.join(self.DATA_FOLDER_PATH, filename + '.' + output_extension)) + new_filename = filename + '.' + output_extension + file_path = os.path.join(self.DATA_FOLDER_PATH, new_filename) + writer.SetFileName(file_path) writer.SetInputConnection(w2if.GetOutputPort()) writer.Write() - # renderer.SetLayer(0) - return + + with open(file_path, "rb") as file: + file_content = file.read() + + return {"blob": self.addAttachment(file_content)} def get_data_base(self): return self.getSharedObject("db") diff --git a/src/tests/conftest.py b/src/tests/conftest.py index 6ab340b..fb91c46 100644 --- a/src/tests/conftest.py +++ b/src/tests/conftest.py @@ -57,7 +57,8 @@ def print_log(self): print(output) def get_response(self): - response = eval(self.ws.recv()) + return self.ws.recv() + response = json.loads(self.ws.recv()) return response def compare_image(self, nb_messages, filename): diff --git a/src/tests/test_protocol.py b/src/tests/test_protocol.py index c2ae3d9..1323058 100644 --- a/src/tests/test_protocol.py +++ b/src/tests/test_protocol.py @@ -103,3 +103,26 @@ def test_set_color(server): server.call("set_color", [{"id": "123456789", "red": 50, "green": 2, "blue": 250}]) assert server.compare_image(3, "set_color.jpeg") == True + + + +def test_take_screenshot(server): + + server.call( + "create_object_pipeline", + [{"id": "123456789", "file_name": "hat.vtp"}], + ) + assert server.compare_image(3, "create_object_pipeline.jpeg") == True + + server.call( + "take_screenshot", + [{"filename": "screenshot", "output_extension": "png", "include_background": False}], + ) + + response1 = server.get_response() + print(f"{response1=}", flush=True) + response2 = server.get_response() + print(f"{response2=}", flush=True) + + blob = response2 + assert type(blob) is bytes diff --git a/src/tests/tests_output/test.jpeg b/src/tests/tests_output/test.jpeg index 3d463c4..6b2509f 100644 Binary files a/src/tests/tests_output/test.jpeg and b/src/tests/tests_output/test.jpeg differ