Add mypy

.github/workflows/pylint.yml

@@ -3,7 +3,7 @@
 #  - Run pylint for each of the supported Python versions. 
 # Pylint will only fail without fixing any of the errors or warnings.
 
-name: Pylint
+name: Pylint & Mypy
 
 on:
   push:
@@ -31,8 +31,11 @@ jobs:
         python -m pip install --upgrade pip
         python -m pip install -r requirement.txt
         python -m pip install -r tests/test-requirement.txt
+        python -m pip install mypy types-click types-protobuf types-redis types-requests types-urllib3 typing_extensions types-cryptography
     - name: Runnning linter.
       run: |
         pylint --load-plugins pylint_quotes --rcfile=.pylintrc --ignore=tests/ agent/
         pylint --load-plugins pylint_quotes --rcfile=.pylintrc -d C0103,W0613 tests/
-
+    - name: Running static types checker.
+      run: |
+        mypy

.github/workflows/pytest.yml

@@ -1,6 +1,6 @@
 # This workflow use pytest:
 #  - Install Python dependencies.
-#  - Run pytest for each of the supported Python versions ["3.8", "3.9", "3.10"]. 
+#  - Run pytest for each of the supported Python versions ["3.10"].
 # Running pytest with -m "no docker" to disable test that require a docker installation.
 
 name: Pytest.
@@ -33,6 +33,7 @@ jobs:
         python -m pip install -r tests/test-requirement.txt
     - name: Runnning tests with pytest.
       run: |
+        set -o pipefail
         pytest -m "not docker" --cov=. | tee pytest-coverage.txt
     - name: Comment coverage
-      uses: coroo/pytest-coverage-commentator@v1.0.2
+      uses: coroo/pytest-coverage-commentator@v1.0.2

.mypy.ini

@@ -0,0 +1,55 @@
+[mypy]
+files = agent, tests
+check_untyped_defs = True
+follow_imports_for_stubs = True
+#disallow_any_decorated = True
+disallow_any_generics = True
+disallow_incomplete_defs = True
+disallow_subclassing_any = True
+disallow_untyped_calls = True
+disallow_untyped_decorators = True
+disallow_untyped_defs = True
+implicit_reexport = False
+no_implicit_optional = True
+show_error_codes = True
+strict_equality = True
+warn_incomplete_stub = True
+warn_redundant_casts = True
+#warn_unreachable = True
+warn_unused_ignores = True
+disallow_any_unimported = True
+warn_return_any = True
+exclude = .*_pb2.py
+
+[mypy-sqlalchemy.*]
+ignore_missing_imports = True
+
+[mypy-docker.*]
+ignore_missing_imports = True
+
+[mypy-markdownify]
+ignore_missing_imports = True
+
+[mypy-jsonschema]
+ignore_missing_imports = True
+
+[mypy-semver]
+ignore_missing_imports = True
+
+[mypy-opentelemetry.*]
+ignore_missing_imports = True
+
+[mypy-ostorlab.agent.*]
+ignore_missing_imports = True
+
+[mypy-ostorlab.runtimes.*]
+ignore_missing_imports = True
+
+[mypy-ipwhois.*]
+ignore_missing_imports = True
+
+[mypy-OpenSSL]
+ignore_missing_imports = True
+
+[mypy-tenacity.*]
+implicit_reexport = True

agent/ip2geo.py

@@ -26,7 +26,7 @@ def __init__(self, endpoint: str = GEO_LOCATION_API_ENDPOINT) -> None:
         """
         self._endpoint = endpoint
 
-    def _locate_ip(self, ip_address) -> Dict[str, Any]:
+    def _locate_ip(self, ip_address: str) -> Any:
         """Get geolocation details of an IP address"""
 
         fields_params = ','.join(GEO_LOCATION_FIELDS)
@@ -39,7 +39,7 @@ def _parse_response(self, response: Dict[str, Any]) -> Dict[str, Any]:
         if response.get('status') == 'success':
             message = {
                 'host': response.get('query'),
-                'version': int(ipaddress.ip_network(response.get('query'), strict=False).version),
+                'version': int(ipaddress.ip_network(str(response.get('query')), strict=False).version),
                 'continent': response.get('continent'),
                 'continent_code': response.get('continentCode'),
                 'country': response.get('country'),
@@ -58,7 +58,7 @@ def _parse_response(self, response: Dict[str, Any]) -> Dict[str, Any]:
             }
         return message
 
-    def get_geolocation_details(self, ip_address: str) -> Dict[str, Any]:
+    def get_geolocation_details(self, ip_address: str) -> None | Dict[str, Any]:
         """Find geolocation details of an IP address.
 
         Args:
@@ -70,3 +70,5 @@ def get_geolocation_details(self, ip_address: str) -> Dict[str, Any]:
         if response is not None:
             geolocation_details = self._parse_response(response)
             return geolocation_details
+        else:
+            return None

agent/request_sender.py

@@ -1,7 +1,7 @@
 """Module responsible for sending HTTP requests"""
 import json
 import logging
-from typing import Dict, Optional
+from typing import Dict, Optional, Any
 
 import requests
 import tenacity
@@ -13,12 +13,13 @@
 class AuthenticationError(Exception):
     """Authentication Error."""
 
+
 @tenacity.retry(
-    stop=tenacity.stop_after_attempt(3),
-    wait=tenacity.wait_fixed(2),
-    retry=tenacity.retry_if_exception_type(),
+    stop=tenacity.stop.stop_after_attempt(3),
+    wait=tenacity.wait.wait_fixed(2),
+    retry=tenacity.retry_if_exception_type(Exception),
     retry_error_callback=lambda retry_state: None)
-def make_request(method: str, path: str, data: Optional[Dict[str, str]] = None):
+def make_request(method: str, path: str, data: Optional[Dict[str, str]] = None) -> Any:
     """Sends an HTTP request.
     Args:
         method: One of HTTP requests, e.g., GET, POST.

agent/utils/ip_range_visitor.py

@@ -1,6 +1,6 @@
 """ Module Responsible for sending ip range geolocation"""
 import ipaddress
-from typing import Callable, Tuple, Any
+from typing import Callable, Tuple, Any, List
 
 from agent import ip2geo
 from agent.ip2geo import logger
@@ -27,7 +27,7 @@ def dichotomy_ip_network_visit(self, ip_network: ipaddress.IPv4Network | ipaddre
         if should_continue is False:
             return
 
-        subnets = list(ip_network.subnets())
+        subnets: List[ipaddress.IPv4Network | ipaddress.IPv6Network] = list(ip_network.subnets())
 
         if len(subnets) == 1:
             # reached the last block.

tests/agent_test.py

@@ -1,11 +1,20 @@
 """Unittests for the Ip2Geo Agent."""
 import re
 import pytest
+from typing import List, Dict, Union
 
 from ostorlab.agent.message import message
+import requests_mock as rq_mock
 
+from agent import ip2geo_agent as ip2geo_agt
 
-def testAgentIp2Geo_whenLocatesIpAddress_emitsBackFindings(ip2geo_agent, agent_mock, agent_persist_mock, requests_mock):
+
+def testAgentIp2Geo_whenLocatesIpAddress_emitsBackFindings(
+        ip2geo_agent: ip2geo_agt.Ip2GeoAgent,
+        agent_mock: List[message.Message],
+        agent_persist_mock: Dict[Union[str, bytes],
+                                 Union[str, bytes]],
+        requests_mock: rq_mock.mocker.Mocker) -> None:
     """Unittest for emitting back the geolocation details found by the Ip2Geo agent."""
     del agent_persist_mock
     matcher = re.compile('http://ip-api.com/json/')
@@ -32,11 +41,11 @@ def testAgentIp2Geo_whenLocatesIpAddress_emitsBackFindings(ip2geo_agent, agent_m
     assert agent_mock[0].data['country_code'] == 'CA'
 
 
-def testAgentIp2Geo_whenIpAddressHasAlreadyBeenProcessed_shouldSkip(ip2geo_agent,
-                                                                    mocker,
-                                                                    agent_mock,
-                                                                    agent_persist_mock,
-                                                                    requests_mock):
+def testAgentIp2Geo_whenIpAddressHasAlreadyBeenProcessed_shouldSkip(ip2geo_agent: ip2geo_agt.Ip2GeoAgent,
+                                                                    agent_mock: List[message.Message],
+                                                                    agent_persist_mock: Dict[Union[str, bytes],
+                                                                                             Union[str, bytes]],
+                                                                    requests_mock: rq_mock.mocker.Mocker) -> None:
     """Unittest for Ip2Geo Agent, when it receives an ip address that has already been processed,
        the agent should skip it.
     """
@@ -61,14 +70,13 @@ def testAgentIp2Geo_whenIpAddressHasAlreadyBeenProcessed_shouldSkip(ip2geo_agent
     ip2geo_agent.process(msg)
 
     assert len(agent_mock) == 256
-    assert mocker is not None
 
 
-def testAgentIp2Geo_whenIpAddressHasHostBits_shouldSkip(ip2geo_agent,
-                                                                    mocker,
-                                                                    agent_mock,
-                                                                    agent_persist_mock,
-                                                                    requests_mock):
+def testAgentIp2Geo_whenIpAddressHasHostBits_shouldSkip(ip2geo_agent: ip2geo_agt.Ip2GeoAgent,
+                                                        agent_mock: List[message.Message],
+                                                        agent_persist_mock: Dict[Union[str, bytes],
+                                                                                 Union[str, bytes]],
+                                                        requests_mock: rq_mock.mocker.Mocker) -> None:
     """Unittest for Ip2Geo Agent, when it receives an ip address that has already been processed,
        the agent should skip it.
     """
@@ -94,4 +102,3 @@ def testAgentIp2Geo_whenIpAddressHasHostBits_shouldSkip(ip2geo_agent,
     ip2geo_agent.process(msg)
 
     assert len(agent_mock) == 256
-    assert mocker is not None

tests/conftest.py

@@ -1,14 +1,17 @@
 """Pytest fixtures for the Ip2Geo agent"""
 import pathlib
+from typing import List
 
 import pytest
+from ostorlab.agent.message import message
 from ostorlab.agent import definitions as agent_definitions
 from ostorlab.runtimes import definitions as runtime_definitions
 
 from agent import ip2geo_agent
 
+
 @pytest.fixture(scope='function', name='ip2geo_agent')
-def fixture_io2geo_agent(agent_mock):
+def fixture_io2geo_agent(agent_mock: List[message.Message]) -> ip2geo_agent.Ip2GeoAgent:
     with (pathlib.Path(__file__).parent.parent / 'ostorlab.yaml').open() as yaml_o:
         definition = agent_definitions.AgentDefinition.from_yaml(yaml_o)
         settings = runtime_definitions.AgentSettings(

tests/utils/test_ip_range_visitor.py

@@ -5,11 +5,12 @@
 
 from agent import ip2geo
 from agent.utils.ip_range_visitor import IpRangeVisitor
+import requests_mock as rq_mock
 
 ip_range_visitor = IpRangeVisitor()
 
 
-def testVistor_withMatchingIPAndMaskRecieved_returnsLocations():
+def testVistor_withMatchingIPAndMaskRecieved_returnsLocations() -> None:
     results = []
     for result in ip_range_visitor.dichotomy_ip_network_visit(ipaddress.ip_network('8.8.8.0/22'),
                                                               ip_range_visitor.is_first_last_ip_same_geolocation):
@@ -79,13 +80,14 @@ def testVistor_withMatchingIPAndMaskRecieved_returnsLocations():
                             'timezone': 'America/Chicago'})]
 
 
-def testVistor_withMaskNotRecieved_returnsIfFirstIPGeolocationEqualLastIPGeolocation():
+def testVistor_withMaskNotRecieved_returnsIfFirstIPGeolocationEqualLastIPGeolocation() -> None:
     for result in ip_range_visitor.dichotomy_ip_network_visit(ipaddress.ip_network('8.8.8.0/32'),
                                                               ip_range_visitor.is_first_last_ip_same_geolocation):
         assert result[0] == result[1]
 
 
-def testIsFirstLastIPSameGeolocation_withNoneLatAndLon_returnsTupleTrueNone(mocker, requests_mock):
+def testIsFirstLastIPSameGeolocation_withNoneLatAndLon_returnsTupleTrueNone(
+        requests_mock: rq_mock.mocker.Mocker) -> None:
     matcher = re.compile('http://ip-api.com/json/')
     requests_mock.get(
         matcher,
@@ -107,12 +109,13 @@ def testIsFirstLastIPSameGeolocation_withNoneLatAndLon_returnsTupleTrueNone(mock
     assert same_geo_location_mocker[1][1]['longitude'] is None
 
 
-def testVistor_withLocatorReturnsNone_returnsIfFirstIPGeolocationEqualLastIPGeolocation(requests_mock):
+def testVistor_withLocatorReturnsNone_returnsIfFirstIPGeolocationEqualLastIPGeolocation(
+        requests_mock: rq_mock.mocker.Mocker) -> None:
     fields_params = ','.join(ip2geo.GEO_LOCATION_FIELDS)
     requests_mock.get(
         f'http://ip-api.com/json//199.102.178.251?fields={fields_params}',
-                                                          status_code=200,
-                                                          json={'errors': 'Too many calls'})
+        status_code=200,
+        json={'errors': 'Too many calls'})
     for result in ip_range_visitor.dichotomy_ip_network_visit(ipaddress.ip_network('199.102.178.251/32'),
                                                               ip_range_visitor.is_first_last_ip_same_geolocation):
         assert result[0] is None