SQ index 404

Author: ybnd

shapeflow/server.py

@@ -5,7 +5,8 @@
 from threading import Thread, Event, Lock
 from typing import Optional
 
-from flask import Flask, send_from_directory, jsonify, request, Response, make_response, abort
+import flask
+from flask import Flask, jsonify, request, Response, make_response, abort
 import waitress
 import webbrowser
 
@@ -106,11 +107,14 @@ def _index():
                     log.error(f"no compiled UI in '{DIST}'")
                 log.error(f"follow the instructions on the 404 page to restore the application")
 
-                return send_from_directory(UI, '404.html')
+                return flask.send_from_directory(UI, '404.html'), 404
 
         @app.route('/<path:file>', methods=['GET'])
         def _get_file(file: str):
-            return self.get_file(file)
+            try:
+                return self.get_file(file)
+            except FileNotFoundError:
+                abort(404)
 
         @app.route('/api/<path:address>', methods=['GET', 'POST', 'PUT'])
         def _call_api(address: str):
@@ -172,7 +176,7 @@ def get_file(self, file: str):
             log.error(f"no such file: '{file}'")
             raise FileNotFoundError
         log.debug(f"serving '{file}'")
-        return send_from_directory(
+        return flask.send_from_directory(
             os.path.dirname(path),
             os.path.basename(path)
         )

test/test_server.py

@@ -1,13 +1,18 @@
 import os
 import shutil
 import unittest
+from unittest.mock import patch
 from contextlib import contextmanager
 
 import time
 import json
 import subprocess
 
+import flask
+from flask.testing import FlaskClient
+
 from shapeflow import settings, ROOTDIR, save_settings
+from shapeflow.server import ShapeflowServer, UI, DIST
 
 CACHE = os.path.join(ROOTDIR, 'test_main-cache')
 DB = os.path.join(ROOTDIR, 'test_main-history.db')
@@ -169,6 +174,55 @@ def test_set_settings_restart(self):
             post(api('set_settings'), data=json.dumps({'settings': first_settings}))
             time.sleep(10)
 
+
+@patch('os.path.isfile')
+@patch('os.path.isdir')
+@patch('flask.send_from_directory')
+class FlaskTest(unittest.TestCase):
+    client: FlaskClient
+
+    def setUp(self) -> None:
+        sfs = ShapeflowServer()
+        sfs._app.config['TESTING'] = True
+
+        self.client = sfs._app.test_client()
+
+    def test_serve_index_200(self, send_from_directory, _, isfile):
+        isfile.return_value = True
+        send_from_directory.return_value = flask.Response()
+
+        r = self.client.get('/')
+
+        send_from_directory.assert_called_with(DIST, 'index.html')
+        self.assertEqual(r.status_code, 200)
+
+    def test_serve_index_404(self, send_from_directory, _, isfile):
+        isfile.return_value = False
+        send_from_directory.return_value = flask.Response()
+
+        r = self.client.get('/')
+
+        send_from_directory.assert_called_with(UI, '404.html')
+        self.assertEqual(r.status_code, 404)
+
+    def test_serve_file_200(self, send_from_directory, _, isfile):
+        isfile.return_value = True
+        send_from_directory.return_value = flask.Response()
+
+        r = self.client.get('/something.txt')
+
+        send_from_directory.assert_called_with(DIST, 'something.txt')
+        self.assertEqual(r.status_code, 200)
+
+    def test_serve_file_404(self, send_from_directory, _, isfile):
+        isfile.return_value = False
+        send_from_directory.return_value = flask.Response()
+
+        r = self.client.get('/something.txt')
+
+        send_from_directory.assert_not_called()
+        self.assertEqual(r.status_code, 404)
+
+
 if __name__ == '__main__':
     unittest.main()
-