From 66866267c5aa03bab02e33d2d3eb32ecb651e4cc Mon Sep 17 00:00:00 2001
From: Alvaro Lopez Garcia <aloga@ifca.unican.es>
Date: Thu, 13 Jul 2023 12:14:30 +0200
Subject: [PATCH] Allow to disable doc endpoint

---
 deepaas/api/__init__.py      |  3 ++-
 deepaas/cmd/run.py           |  2 +-
 deepaas/config.py            |  8 ++++++++
 deepaas/tests/test_v2_api.py | 27 +++++++++++++++++++++++++++
 4 files changed, 38 insertions(+), 2 deletions(-)

diff --git a/deepaas/api/__init__.py b/deepaas/api/__init__.py
index b4db5aee..eb1cd738 100644
--- a/deepaas/api/__init__.py
+++ b/deepaas/api/__init__.py
@@ -54,6 +54,7 @@
 
 async def get_app(
     swagger=True,
+    enable_doc=True,
     doc="/ui",
     prefix="",
     static_path="/static/swagger",
@@ -101,7 +102,7 @@ async def get_app(
             basePath=base_path,
             version=deepaas.__version__,
             url="/swagger.json",
-            swagger_path=doc if doc else None,
+            swagger_path=doc if enable_doc else None,
             prefix=prefix,
             static_path=static_path,
             in_place=True,
diff --git a/deepaas/cmd/run.py b/deepaas/cmd/run.py
index 3989dbb1..b68e46b8 100644
--- a/deepaas/cmd/run.py
+++ b/deepaas/cmd/run.py
@@ -92,7 +92,7 @@ def main():
     log.info("Starting DEEPaaS version %s", deepaas.__version__)
 
     app = api.get_app(
-        doc="/ui",
+        enable_doc=CONF.doc_endpoint,
         enable_train=CONF.train_endpoint,
         enable_predict=CONF.predict_endpoint,
     )
diff --git a/deepaas/config.py b/deepaas/config.py
index fcd4e494..18b079a0 100644
--- a/deepaas/config.py
+++ b/deepaas/config.py
@@ -48,6 +48,14 @@
 print to the standard output and error (i.e. stdout and stderr) through the
 "/debug" endpoint. Default is to not provide this information. This will not
 provide logging information about the API itself.
+""",
+    ),
+    cfg.BoolOpt(
+        "doc-endpoint",
+        default=True,
+        help="""
+Enable documentation endpoint. If set we will provide the documentation
+through the "/ui" endpoint. Default is to provide this information.
 """,
     ),
     cfg.IntOpt(
diff --git a/deepaas/tests/test_v2_api.py b/deepaas/tests/test_v2_api.py
index 482c35ed..e7d90dc5 100644
--- a/deepaas/tests/test_v2_api.py
+++ b/deepaas/tests/test_v2_api.py
@@ -137,6 +137,33 @@ async def test_get_metadata(self):
         self.assertDictEqual(meta["models"][0], await ret.json())
 
 
+class TestApiV2NoDoc(base.TestCase):
+    async def get_application(self):
+        app = web.Application(debug=True)
+        app.middlewares.append(web.normalize_path_middleware())
+
+        deepaas.model.v2.register_models(app)
+
+        v2app = v2.get_app(enable_doc=False)
+        app.add_subapp("/v2", v2app)
+
+        return app
+
+    def setUp(self):
+        super(TestApiV2NoDoc, self).setUp()
+
+        self.maxDiff = None
+
+        self.flags(debug=True)
+
+    def assert_ok(self, response):
+        self.assertIn(response.status, [200, 201])
+
+    async def test_not_found(self):
+        ret = await self.client.get("/ui")
+        self.assertEqual(404, ret.status)
+
+
 class TestApiV2(base.TestCase):
     async def get_application(self):
         app = web.Application(debug=True)