Added typing and fixed unit tests

oravidov · oravidov · commit 05e98be6de73 · 2021-10-03T10:00:13.000+03:00
diff --git a/lineage/query_context.py b/lineage/query_context.py
@@ -1,12 +1,13 @@
 from datetime import datetime
-from typing import Optional
+from typing import Optional, Union
 import dateutil.parser
 
 
 class QueryContext(object):
-    def __init__(self, queried_database: str, queried_schema: str, query_time: Optional[datetime] = None,
-                 query_volume: Optional[int] = None, query_type: Optional[str] = None, user_name: Optional[str] = None,
-                 role_name: Optional[str] = None):
+    def __init__(self, queried_database: Optional[str] = None, queried_schema: Optional[str] = None,
+                 query_time: Optional[datetime] = None, query_volume: Optional[int] = None,
+                 query_type: Optional[str] = None, user_name: Optional[str] = None,
+                 role_name: Optional[str] = None) -> None:
         self.queried_database = queried_database
         self.queried_schema = queried_schema
         self.query_time = query_time
@@ -15,32 +16,38 @@ def __init__(self, queried_database: str, queried_schema: str, query_time: Optio
         self.user_name = user_name
         self.role_name = role_name
 
-    def to_dict(self):
-        query_time_str = self.query_time.isoformat() if self.query_time is not None else None
+    def to_dict(self) -> dict:
         return {'queried_database': self.queried_database,
                 'queried_schema': self.queried_schema,
-                'query_time': query_time_str,
+                'query_time': self._query_time_to_str(self.query_time),
                 'query_volume': self.query_volume,
                 'query_type': self.query_type,
                 'user_name': self.user_name,
                 'role_name': self.role_name}
 
     @staticmethod
-    def _html_param_with_default(param, default='unknown'):
+    def _query_time_to_str(query_time: Optional[datetime], fmt: str = None) -> Optional[str]:
+        if query_time is None:
+            return None
+
+        if fmt is None:
+            return query_time.isoformat()
+
+        return query_time.strftime(fmt)
+
+    @staticmethod
+    def _html_param_with_default(param: Union[str, int], default: Union[str, int] = 'unknown') -> Union[str, int]:
         return default if param is None else param
 
-    def to_html(self):
+    def to_html(self) -> str:
         query_type = self._html_param_with_default(self.query_type)
         user_name = self._html_param_with_default(self.user_name)
         role_name = self._html_param_with_default(self.role_name)
-        query_time = self.query_time.strftime('%Y-%m-%d %H:%M:%S')
-
-        if self.query_volume is not None and self.query_volume > 0:
-            volume_color = "DarkSlateGrey"
-            query_volume = self.query_volume
-        else:
+        query_time = self._query_time_to_str(self.query_time, fmt='%Y-%m-%d %H:%M:%S')
+        query_volume = self._html_param_with_default(self.query_volume, 0)
+        volume_color = "DarkSlateGrey"
+        if query_volume == 0:
             volume_color = "tomato"
-            query_volume = 0
 
         return f"""
                     <html>
@@ -60,7 +67,7 @@ def to_html(self):
         """
 
     @staticmethod
-    def from_dict(query_context_dict):
+    def from_dict(query_context_dict: dict) -> 'QueryContext':
         if 'query_time' in query_context_dict and query_context_dict['query_time'] is not None:
             query_context_dict['query_time'] = dateutil.parser.parse(query_context_dict['query_time'])
         return QueryContext(**query_context_dict)
diff --git a/tests/test_lineage_graph.py b/tests/test_lineage_graph.py
@@ -75,13 +75,14 @@ def test_lineage_graph_add_nodes_and_edges(sources, targets, edges, show_isolate
     di_graph_mock = mock.create_autospec(nx.DiGraph)
     reference._lineage_graph = di_graph_mock
 
-    reference._add_nodes_and_edges(sources, targets)
+    empty_query_context = QueryContext()
+    reference._add_nodes_and_edges(sources, targets, empty_query_context)
 
     node_calls = []
     if len(sources) > 0:
         node_calls.append(mock.call(sources))
     if len(targets) > 0:
-        node_calls.append(mock.call(targets))
+        node_calls.append(mock.call(targets, title=empty_query_context.to_html()))
 
     edge_calls = []
     for edge in edges:
