From 648cf4c76fae1460fdbcfd041779ce9b7ba2a538 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?F=C3=A9lix-Antoine=20Fortin?=
 <felix-antoine.fortin@calculquebec.ca>
Date: Thu, 28 Mar 2019 12:06:25 -0400
Subject: [PATCH] Replace current_ip and current_port by ip and port

---
 batchspawner/api.py          |  9 +++++----
 batchspawner/batchspawner.py | 23 ++++++++++-------------
 2 files changed, 15 insertions(+), 17 deletions(-)

diff --git a/batchspawner/api.py b/batchspawner/api.py
index f78693a4..109e3ba7 100644
--- a/batchspawner/api.py
+++ b/batchspawner/api.py
@@ -5,12 +5,13 @@
 class BatchSpawnerAPIHandler(APIHandler):
     @web.authenticated
     def post(self):
-        """POST set user's spawner port number"""
+        """POST set user spawner data"""
         user = self.get_current_user()
         data = self.get_json_body()
-        port = int(data.get('port', 0))
-        user.spawner.current_port = port
-        self.finish(json.dumps({"message": "BatchSpawner port configured"}))
+        for key, value in data.items():
+            if hasattr(user.spawner, key):
+                setattr(user.spawner, key, value)
+        self.finish(json.dumps({"message": "BatchSpawner data configured"}))
         self.set_status(201)
 
 default_handlers.append((r"/api/batchspawner", BatchSpawnerAPIHandler))
diff --git a/batchspawner/batchspawner.py b/batchspawner/batchspawner.py
index 1a8a77b5..885e020f 100644
--- a/batchspawner/batchspawner.py
+++ b/batchspawner/batchspawner.py
@@ -165,12 +165,6 @@ def _req_keepvars_default(self):
     # Will get the raw output of the job status command unless overridden
     job_status = Unicode()
 
-    # Will get the address of the server as reported by job manager
-    current_ip = Unicode()
-
-    # Will get the port of the server as reported by singleserver
-    current_port = Integer()
-
     # Prepare substitution variables for templates using req_xyz traits
     def get_req_subvars(self):
         reqlist = [ t for t in self.trait_names() if t.startswith('req_') ]
@@ -349,6 +343,9 @@ def poll(self):
     @gen.coroutine
     def start(self):
         """Start the process"""
+        self.ip = self.traits()['ip'].default_value
+        self.port = self.traits()['port'].default_value
+
         if jupyterhub.version_info >= (0,8) and self.server:
             self.server.port = self.port
 
@@ -375,20 +372,20 @@ def start(self):
                            ' after starting.')
             yield gen.sleep(self.startup_poll_interval)
 
-        self.current_ip = self.state_gethost()
-        while self.current_port == 0:
+        self.ip = self.state_gethost()
+        while self.port == 0:
             yield gen.sleep(self.startup_poll_interval)
 
         if jupyterhub.version_info < (0,7):
             # store on user for pre-jupyterhub-0.7:
-            self.user.server.port = self.current_port
-            self.user.server.ip = self.current_ip
+            self.user.server.port = self.port
+            self.user.server.ip = self.ip
         self.db.commit()
         self.log.info("Notebook server job {0} started at {1}:{2}".format(
-                        self.job_id, self.current_ip, self.current_port)
+                        self.job_id, self.ip, self.port)
             )
 
-        return self.current_ip, self.current_port
+        return self.ip, self.port
 
     @gen.coroutine
     def stop(self, now=False):
@@ -408,7 +405,7 @@ def stop(self, now=False):
             yield gen.sleep(1.0)
         if self.job_id:
             self.log.warn("Notebook server job {0} at {1}:{2} possibly failed to terminate".format(
-                             self.job_id, self.current_ip, self.port)
+                             self.job_id, self.ip, self.port)
                 )