nitb_netreg_mass: Allow a suite to configure num subscriber

Remove hardcoded configuration and replace it with a reservation in
the suite.

Change-Id: Ic9fd560c2924731e4fd6eea5aaf3ad565cb4ef52

7 files changed, 80 insertions, 65 deletions

diff --git a/src/osmo_gsm_tester/ms_driver.py b/src/osmo_gsm_tester/ms_driver.py
index e8e543c..8ee93b7 100644
--- a/src/osmo_gsm_tester/ms_driver.py
+++ b/src/osmo_gsm_tester/ms_driver.py
@@ -27,29 +27,6 @@ import os.path
 import shutil
 import tempfile
 
-class Subscriber(log.Origin):
-    def __init__(self, imsi, ki):
-        super().__init__(log.C_RUN, 'subscriber')
-        self._imsi = imsi
-        self._ki = ki
-        self._auth_algo = "comp128v1"
-        self._msisdn = None
-
-    def msisdn(self):
-        return self._msisdn
-
-    def set_msisdn(self, msisdn):
-        self._msisdn = msisdn
-
-    def imsi(self):
-       return self._imsi
-
-    def ki(self):
-       return self._ki
-
-    def auth_algo(self):
-       return self._auth_algo
-
 class MsDriver(log.Origin):
 
     def __init__(self, suite_run):
@@ -57,7 +34,6 @@ class MsDriver(log.Origin):
         self._suite_run = suite_run
 
         # TODO: take config out of the test scenario
-        self._num_ms = 10
         self._time_start = timedelta(seconds=60)
         self._time_step = timedelta(milliseconds=100)
         self._test_duration = timedelta(seconds=120)
@@ -65,6 +41,8 @@ class MsDriver(log.Origin):
         self._loop = SimpleLoop()
         self._test_case = None
         self.event_server_sk_tmp_dir = None
+        self._subscribers = []
+        self._configured = False
 
         if len(self.event_server_path().encode()) > 107:
             raise log.Error('Path for event_server socket is longer than max allowed len for unix socket path (107):', self.event_server_path())
@@ -101,6 +79,10 @@ class MsDriver(log.Origin):
         mobile = check_and_return_binary('mobile')
         return BinaryOptions(virtphy, mobile, env)
 
+    def subscriber_add(self, subscriber):
+        """Adds a subscriber to the list of subscribers."""
+        self._subscribers.append(subscriber)
+
     def configure(self):
         """
         Configures the subscribers, tests and registration server. Needs to be
@@ -111,24 +93,14 @@ class MsDriver(log.Origin):
         self._ev_server = EventServer("ev_server", event_server_path)
         self._ev_server.listen(self._loop)
         options = self.build_binary_options()
-        self._test_case = MassUpdateLocationTest("mass", options, self._num_ms, self._cdf,
+        self._test_case = MassUpdateLocationTest("mass", options, self._cdf,
                                                  self._ev_server,
                                                  util.Dir(self.event_server_sk_tmp_dir),
                                                  suite_run=self._suite_run)
 
-        # TODO: We should pass subscribers down to the test and not get it from
-        # there.
-        self._subs = [Subscriber(imsi=mob.imsi(), ki=mob.ki()) for mob in self._test_case.mobiles()]
-
-
-    def ms_subscribers(self):
-        """
-        Returns a list of 'subscribers' that were configured in the
-        current scenario.
-        """
-        if not hasattr(self, '_subs'):
-            self.configure()
-        return self._subs
+        for sub in self._subscribers:
+            self._test_case.subscriber_add(sub)
+        self._configured = True
 
     def run_test(self):
         """
@@ -136,7 +108,7 @@ class MsDriver(log.Origin):
         devices according to their schedule. Returns once all tests succeeded
         or the configured timeout has passed.
         """
-        if not hasattr(self, '_subs'):
+        if not self._configured:
             self.configure()
         self._test_case.run_test(self._loop, self._test_duration)
 
diff --git a/src/osmo_gsm_tester/ms_osmo_mobile.py b/src/osmo_gsm_tester/ms_osmo_mobile.py
index 1c28117..f497772 100644
--- a/src/osmo_gsm_tester/ms_osmo_mobile.py
+++ b/src/osmo_gsm_tester/ms_osmo_mobile.py
@@ -26,3 +26,9 @@ class MSOsmoMobile(ms.MS):
     def cleanup(self):
         # do nothing for a virtual resource
         pass
+
+    def ki(self):
+        ki = super().ki()
+        if not ki:
+            return "00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00"
+        return ki
diff --git a/src/osmo_gsm_tester/suite.py b/src/osmo_gsm_tester/suite.py
index db580ec..aefdc24 100644
--- a/src/osmo_gsm_tester/suite.py
+++ b/src/osmo_gsm_tester/suite.py
@@ -344,6 +344,15 @@ class SuiteRun(log.Origin):
             l.append(self.modem())
         return l
 
+    def all_resources(self, resource_func):
+        """Returns all yielded resource."""
+        l = []
+        while True:
+            try:
+                l.append(resource_func())
+            except resource.NoResourceExn:
+                return l
+
     def esme(self):
         esme_obj = esme.Esme(self.msisdn())
         self.register_for_cleanup(esme_obj)
diff --git a/src/osmo_ms_driver/__main__.py b/src/osmo_ms_driver/__main__.py
index 7e1afa8..24acd61 100644
--- a/src/osmo_ms_driver/__main__.py
+++ b/src/osmo_ms_driver/__main__.py
@@ -21,7 +21,9 @@ from .simple_loop import SimpleLoop
 from .location_update_test import MassUpdateLocationTest
 from .cdf import cdfs
 from .starter import BinaryOptions
+from .test_support import imsi_ki_gen
 from osmo_gsm_tester import log, util
+from osmo_gsm_tester import ms
 
 # System modules
 from datetime import timedelta
@@ -84,7 +86,19 @@ def main():
 
     # Just a single test for now.
     options = BinaryOptions("virtphy", "mobile", os.environ)
-    test = MassUpdateLocationTest("lu_test", options, args.num_ms, cdf, ev_server, tmp_dir)
+    test = MassUpdateLocationTest("lu_test", options, cdf, ev_server, tmp_dir)
+
+    # Add subscribers to the test.
+    imsi_gen = imsi_ki_gen()
+    for i in range(0, args.num_ms):
+        imsi, ki = next(imsi_gen)
+        conf = {
+            'imsi': imsi,
+            'ki': ki,
+            'auth_algo': 'comp128v1',
+        }
+        test.subscriber_add(ms.MS("ms_%d" % i, conf))
+
     atexit.register(test.stop_all)
 
     # Run until everything has been launched
diff --git a/src/osmo_ms_driver/location_update_test.py b/src/osmo_ms_driver/location_update_test.py
index 8500607..2e3a648 100644
--- a/src/osmo_ms_driver/location_update_test.py
+++ b/src/osmo_ms_driver/location_update_test.py
@@ -19,7 +19,7 @@
 from copy import copy
 from osmo_gsm_tester import log
 from .starter import OsmoVirtPhy, OsmoMobile
-from .test_support import imsi_ki_gen, Results
+from .test_support import Results
 
 from datetime import timedelta
 
@@ -56,42 +56,53 @@ class MassUpdateLocationTest(log.Origin):
     TEMPLATE_LUA = "osmo-mobile-lu.lua"
     TEMPLATE_CFG = "osmo-mobile.cfg"
 
-    def __init__(self, name, options, number_of_ms, cdf_function,
+    def __init__(self, name, options, cdf_function,
                  event_server, tmp_dir, suite_run=None):
         super().__init__(log.C_RUN, name)
         self._binary_options = options
-        self._number_of_ms = number_of_ms
         self._cdf = cdf_function
-        self._cdf.set_target(number_of_ms)
         self._suite_run = suite_run
+        self._tmp_dir = tmp_dir
         self._unstarted = []
         self._mobiles = []
         self._phys = []
         self._results = {}
-        imsi_gen = imsi_ki_gen()
 
-        self._outstanding = number_of_ms
-        for i in range(0, number_of_ms):
+        self._event_server = event_server
+        self._event_server.register(self.handle_msg)
+        self._started = []
+        self._subscribers = []
+
+    def subscriber_add(self, subscriber):
+        """
+        Adds a subscriber to the list of subscribers.
+
+        Must be called before starting the testcase.
+        """
+        self._subscribers.append(subscriber)
+
+    def configure_tasks(self):
+        """Sets up the test run."""
+
+        self._cdf.set_target(len(self._subscribers))
+        self._outstanding = len(self._subscribers)
+        for i in range(0, self._outstanding):
             ms_name = "%.5d" % i
 
-            phy = OsmoVirtPhy(options.virtphy, options.env,
-                              ms_name, tmp_dir)
+            phy = OsmoVirtPhy(self._binary_options.virtphy,
+                              self._binary_options.env,
+                              ms_name, self._tmp_dir)
             self._phys.append(phy)
 
-            launcher = OsmoMobile(options.mobile, options.env,
-                                ms_name, tmp_dir, self.TEMPLATE_LUA,
-                                self.TEMPLATE_CFG, imsi_gen,
+            launcher = OsmoMobile(self._binary_options.mobile,
+                                self._binary_options.env,
+                                ms_name, self._tmp_dir, self.TEMPLATE_LUA,
+                                self.TEMPLATE_CFG, self._subscribers[i],
                                 phy.phy_filename(),
-                                event_server.server_path())
+                                self._event_server.server_path())
             self._results[ms_name] = LUResult(ms_name)
             self._mobiles.append(launcher)
-        self._event_server = event_server
-        self._event_server.register(self.handle_msg)
         self._unstarted = copy(self._mobiles)
-        self._started = []
-
-    def mobiles(self):
-        return self._mobiles
 
     def pre_launch(self, loop):
         """
@@ -111,6 +122,7 @@ class MassUpdateLocationTest(log.Origin):
     def prepare(self, loop):
         self.log("Starting testcase")
 
+        self.configure_tasks()
         self.pre_launch(loop)
 
         self._start_time = time.clock_gettime(time.CLOCK_MONOTONIC)
diff --git a/src/osmo_ms_driver/starter.py b/src/osmo_ms_driver/starter.py
index dd3aac5..010947b 100644
--- a/src/osmo_ms_driver/starter.py
+++ b/src/osmo_ms_driver/starter.py
@@ -69,14 +69,14 @@ class OsmoVirtPhy(Launcher):
             self._vphy_proc.terminate()
 
 class OsmoMobile(Launcher):
-    def __init__(self, binary, env, name_number, tmp_dir, lua_tmpl, cfg_tmpl, imsi_ki_generator, phy_filename, ev_server_path):
+    def __init__(self, binary, env, name_number, tmp_dir, lua_tmpl, cfg_tmpl, subscriber, phy_filename, ev_server_path):
         super().__init__(binary, env, "osmo-ms-mob", name_number, tmp_dir)
         self._lua_template = lua_tmpl
         self._cfg_template = cfg_tmpl
-        self._imsi_ki_generator = imsi_ki_generator
         self._phy_filename = phy_filename
         self._ev_server_path = ev_server_path
-        self._imsi, self._ki = next(self._imsi_ki_generator)
+        self._imsi = subscriber.imsi()
+        self._ki = subscriber.ki()
         self._omob_proc = None
 
     def imsi(self):
diff --git a/suites/nitb_netreg_mass/register_default_mass.py b/suites/nitb_netreg_mass/register_default_mass.py
index ede2b74..e7c91da 100644
--- a/suites/nitb_netreg_mass/register_default_mass.py
+++ b/suites/nitb_netreg_mass/register_default_mass.py
@@ -9,6 +9,7 @@ print('use resources...')
 nitb = suite.nitb()
 bts = suite.bts()
 ms_driver = suite.ms_driver()
+modems = suite.all_resources(suite.modem)
 
 print('start nitb and bts...')
 nitb.bts_add(bts)
@@ -16,9 +17,10 @@ nitb.start()
 bts.start()
 wait(nitb.bts_is_connected, bts)
 
-# Configure all MS that the MS driver knows about.
-for ms in ms_driver.ms_subscribers():
-    nitb.subscriber_add(ms)
+# Configure all MS that are available to this test.
+for modem in modems:
+    nitb.subscriber_add(modem)
+    ms_driver.subscriber_add(modem)
 
 # Run the base test.
 ms_driver.run_test()