magma · nstng · Mar 24, 2022 · Mar 11, 2022 · Mar 11, 2022 · Mar 10, 2022
@@ -13,19 +13,10 @@
 import os
 import subprocess  # noqa: S404
 from pathlib import Path
-from typing import List
+from typing import List, Optional
 
 from lte.protos.subscriberdb_pb2 import SubscriberID
 
-parents = Path.cwd().parents
-parts = Path.cwd().parts
-home = str(Path.home())
-if 'lte' in parts and len(parents) > 3:
-    # Get relative import path for protos
-    IMPORT_PATH = parents[3]
-else:
-    IMPORT_PATH = str(home) + '/magma'
-
 RESULTS_PATH = '/var/tmp'
 PROTO_DIR = 'lte/protos'
 
@@ -84,7 +75,7 @@ def benchmark_grpc_request(
     output_file: str,
     num_reqs: int,
     address: str,
-    import_path: str = None,
+    import_path: str,
 ):
     """Run GHZ based GRPC benchmarking
 
@@ -95,16 +86,22 @@ def benchmark_grpc_request(
         output_file (str): a path where result is written to
         num_reqs (int): number of requests to send
         address (str): address to the service being benchmarked
+        import_path (str): protobuf import path
     """
-    import_path = import_path or IMPORT_PATH
     if not Path(import_path).exists():
         print('Protobuf import path directory does not exist, exiting')
         return
     cmd_list = [
         'ghz',
-        '--insecure', '--proto', proto_path, '-i', import_path, '--total',
-        str(num_reqs), '--call', full_request_type, '-D', input_file,
-        '-O', 'json', '-o', output_file, address,
+        '--insecure',
+        '--proto', proto_path,
+        '-i', import_path,
+        '--total', str(num_reqs),
+        '--call', full_request_type,
+        '-D', input_file,
+        '-O', 'json',
+        '-o', output_file,
+        address,
     ]
 
     try:

@@ -17,6 +17,7 @@
 from functools import lru_cache
 from socket import AF_INET, AF_INET6, inet_ntop
 from struct import pack
+from typing import Tuple
 
 import psutil
 from magma.kernsnoopd import metrics
@@ -92,7 +93,7 @@ def _get_cmdline(self, pid: int) -> list:
         return psutil.Process(pid=pid).cmdline()
 
     @lru_cache(maxsize=1024)
-    def _ip_addr_to_str(self, family: int, daddr: (int, int)) -> str:
+    def _ip_addr_to_str(self, family: int, daddr: Tuple[int, int]) -> str:
         """
         _ip_addr_to_str returns a string representation of an IPv4 or IPv6
         address. It caches results in an LRU cache to reduce cost of conversion
@@ -110,6 +111,8 @@ def _ip_addr_to_str(self, family: int, daddr: (int, int)) -> str:
         elif family == AF_INET6:
             # noinspection PyTypeChecker
             return inet_ntop(AF_INET6, self.Addr(*daddr))
+        else:
+            raise Exception("No valid socket family given!")
 
     def handle(self, bpf):
         """

@@ -12,11 +12,12 @@
 """
 import logging
 import os
+from typing import List
 
 from bcc import BPF  # pylint:disable=import-error
 from jinja2 import Template
 from magma.common.job import Job
-from magma.kernsnoopd.handlers import ebpf_handlers
+from magma.kernsnoopd.handlers import ByteCounter, ebpf_handlers
 
 EBPF_SRC_DIR = "/var/opt/magma/ebpf/kernsnoopd/"
 if not os.path.isdir(EBPF_SRC_DIR):
@@ -65,7 +66,7 @@ def __init__(
 
         super().__init__(interval=collect_interval, loop=service_loop)
         self._bpf = None
-        self._handlers = []
+        self._handlers: List[ByteCounter] = []
         self._loop = service_loop
         self._ebpf_programs = programs
         self._service_registry = service_registry

@@ -16,7 +16,11 @@
 import subprocess  # noqa: S404 ignore security warning about subprocess
 from typing import List
 
-MAGMA_ROOT = os.getenv('MAGMA_ROOT')
+if os.getenv('MAGMA_ROOT'):
+    MAGMA_ROOT = os.environ["MAGMA_ROOT"]
+else:
+    raise Exception("'MAGMA_ROOT' needs to be set and point to the Magma root directory!")
+
 LINT_DOCKER_PATH = os.path.join(
     MAGMA_ROOT,
     'lte/gateway/docker/python-precommit/',
@@ -29,9 +33,6 @@
 
 def main() -> None:
     """Provide command-line options to format/lint Magma's Python codebase"""
-    if MAGMA_ROOT is None:
-        print("Please set the 'MAGMA_ROOT' environment variable to point to the root directory")
-        return
     print("Magma root is " + MAGMA_ROOT)
     args = _parse_args()
 

@@ -40,7 +40,7 @@
 DEFAULT_DNS_IPV6_ADDR = "2001:4860:4860:0:0:0:0:8888"
 DEFAULT_P_CSCF_IPV4_ADDR = "172.27.23.150"
 DEFAULT_P_CSCF_IPV6_ADDR = "2a12:577:9941:f99c:0002:0001:c731:f114"
-DEFAULT_NGAP_S_NSSAI_SST = 1
+DEFAULT_NGAP_S_NSSAI_SST = "1"
 DEFAULT_NGAP_S_NSSAI_SD = "ffffff"
 DEFAULT_NGAP_AMF_NAME = "MAGMAAMF1"
 DEFAULT_NGAP_AMF_REGION_ID = "1"

@@ -76,9 +76,9 @@ def _deserialize_generic_json(
                 return element
 
     if isinstance(element, dict):
-        keys = element.keys()
+        keys = list(element.keys())
     elif isinstance(element, list):
-        keys = range(len(element))
+        keys = list(range(len(element)))
     else:
         # in case it is neither of the know elements, just return as is
         return element

@@ -18,7 +18,7 @@
 # We can use an environment variable to pass in the package version during
 # build. Since we don't distribute this on its own, we don't really care what
 # version this represents. 'None' defaults to 0.0.0.
-VERSION = os.environ.get('PKG_VERSION', None)
+VERSION = os.environ.get('PKG_VERSION', '0.0.0')
 
 setup(
     name='lte',

diff --git a/mypy.ini b/mypy.ini
@@ -8,16 +8,9 @@ namespace_packages = True
 install_types = True
 non_interactive = True
 exclude = (?x)(
+          ^lte/gateway/python/integ_tests/ |
           ^lte/gateway/python/magma/pipelined/ |
           ^lte/gateway/python/magma/enodebd/ |
-          ^lte/gateway/python/magma/subscriberdb/protocols/diameter/avp.py$ |
           ^lte/gateway/python/magma/subscriberdb/protocols/m5g_auth_servicer.py$ |
-          ^lte/gateway/python/magma/kernsnoopd/ |
-          ^lte/gateway/python/precommit.py$ |
-          ^lte/gateway/python/scripts/agw_health_cli.py$ |
-          ^lte/gateway/python/load_tests.common.py$ |
-          ^lte/gateway/python/scripts/generate_oai_config.py$ |
-          ^lte/gateway/python/setup.py$ |
-          ^lte/gateway/python/scripts/state_cli.py$ |
-          ^lte/gateway/python/integ_tests/
+          ^lte/gateway/python/scripts/generate_oai_config.py$
           )