diff --git a/src/webui/Config.py b/src/webui/Config.py
index 5ff3e16486c9f94340867d95eb95df77aa4215db..152a4d7c41c5317a64a055edd49ea1c63e6069a5 100644
--- a/src/webui/Config.py
+++ b/src/webui/Config.py
@@ -15,10 +15,10 @@ HOST = '0.0.0.0'  # accepts connections coming from any ADDRESS
 
 DEBUG=False
 
-CONTEXT_SERVICE_ADDRESS = '10.107.242.226'
-# CONTEXT_SERVICE_ADDRESS = 'context'  # TODO: use the correct context ADDRESS
+# CONTEXT_SERVICE_ADDRESS = '10.107.242.226'
+CONTEXT_SERVICE_ADDRESS = 'context'  # TODO: use the correct context ADDRESS
 CONTEXT_SERVICE_PORT = 1010
 
-DEVICE_SERVICE_ADDRESS = '10.109.49.149'
-# DEVICE_SERVICE_ADDRESS = 'device'  # TODO: use the correct address
+# DEVICE_SERVICE_ADDRESS = '10.109.49.149'
+DEVICE_SERVICE_ADDRESS = 'device'  # TODO: use the correct address
 DEVICE_SERVICE_PORT = 2020
diff --git a/src/webui/Dockerfile b/src/webui/Dockerfile
index de5d393c3d875fd485c5464ada395a69e1f2af2b..b65889025c8d52c42ded40fcff0f574a0923d87d 100644
--- a/src/webui/Dockerfile
+++ b/src/webui/Dockerfile
@@ -29,7 +29,12 @@ RUN python3 -m pip install -r webui/requirements.txt
 
 # Add files into working directory
 COPY common/. common
-COPY context/. context
+COPY context/__init__.py context/__init__.py
+COPY context/proto/. context/proto
+COPY context/client/. context/client
+COPY device/__init__.py device/__init__.py
+COPY device/proto/. device/proto
+COPY device/client/. device/client
 COPY webui/. webui
 
 # Start webui service
diff --git a/src/webui/service/device/routes.py b/src/webui/service/device/routes.py
index 080035f91a5c3b69cd5650eec185187fddd06321..30ca43a3af3df695d976b670f1feb3229313d563 100644
--- a/src/webui/service/device/routes.py
+++ b/src/webui/service/device/routes.py
@@ -1,11 +1,11 @@
 from flask import render_template, Blueprint, flash, session, redirect
-from context.proto.context_pb2 import ConfigActionEnum, ConfigRule, TopologyIdList, TopologyList
 from device.client.DeviceClient import DeviceClient
+from context.client.ContextClient import ContextClient
 from webui.Config import (CONTEXT_SERVICE_ADDRESS, CONTEXT_SERVICE_PORT,
                 DEVICE_SERVICE_ADDRESS, DEVICE_SERVICE_PORT)
-from context.client.ContextClient import ContextClient
 from webui.proto.context_pb2 import (ContextId, DeviceList, DeviceId,
-    Device, DeviceDriverEnum, DeviceOperationalStatusEnum)
+    Device, DeviceDriverEnum, DeviceOperationalStatusEnum,
+    ConfigActionEnum, ConfigRule, TopologyIdList, TopologyList)
 from webui.service.device.forms import AddDeviceForm
 
 device = Blueprint('device', __name__, url_prefix='/device')
diff --git a/src/webui/tests/test_unitary.py b/src/webui/tests/test_unitary.py
index 833f30439c3792da8982c6a96b29b814eb2c7197..0c0b82a6bbb096dc825791944331314093236970 100644
--- a/src/webui/tests/test_unitary.py
+++ b/src/webui/tests/test_unitary.py
@@ -95,6 +95,4 @@ def test_device_add_action(client):
         'device_endpoints': [],
     }
     rw = client.post('/device/add', data=DEVICE_EMU, follow_redirects=True)
-    with open('device_add.html', 'wb') as file:
-        file.write(rw.data)
     assert b'success' in rw.data