Add simple (WIP) tests to python code

Author: yujiyokoo

languages/python/jupyter_notebook/cs_models.py

@@ -2,6 +2,7 @@
 from sqlalchemy.types import TypeDecorator, String, Integer, Date, Boolean, Float
 from sqlalchemy import create_engine, select, text
 from sqlalchemy.exc import IntegrityError
+from datetime import datetime
 import json
 
 import sys
@@ -23,6 +24,7 @@ def process_bind_param(self, value, dialect):
                     "c": self.column_name
                 },
                 "v": 1,
+                "q": None
             }
             value = json.dumps(value_dict)
         return value
@@ -35,43 +37,51 @@ def process_result_value(self, value, dialect):
 class EncryptedInt(CsTypeDecorator):
     impl = String
 
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
+    def process_result_value(self, value, dialect):
+        if value is None:
+            return None
+        return int(value['p'])
+
 
 class EncryptedBoolean(CsTypeDecorator):
     impl = String
 
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-
     def process_bind_param(self, value, dialect):
         if value is not None:
             value = str(value).lower()
         return super().process_bind_param(value, dialect)
 
+    def process_result_value(self, value, dialect):
+        if value is None:
+            return None
+        return value['p'] == 'true'
+
 class EncryptedDate(CsTypeDecorator):
     impl = String
 
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
+    def process_result_value(self, value, dialect):
+        if value is None:
+            return None
+        return datetime.fromisoformat(value['p']).date()
 
 class EncryptedFloat(CsTypeDecorator):
     impl = String
 
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
+    def process_result_value(self, value, dialect):
+        if value is None:
+            return None
+        return float(value['p'])
 
 class EncryptedUtf8Str(CsTypeDecorator):
     impl = String
 
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-
 class EncryptedJsonb(CsTypeDecorator):
     impl = String
 
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
+    def process_result_value(self, value, dialect):
+        if value is None:
+            return None
+        return json.loads(value['p'])
 
 class BaseModel(DeclarativeBase):
     pass

languages/python/jupyter_notebook/cs_models_test.py

@@ -0,0 +1,46 @@
+import unittest
+from datetime import date
+
+from cs_models import *
+
+class TestExampleModel(unittest.TestCase):
+    def setUp(self):
+        # TODO: configure database URL in environment variable and use a test db (not getting_started)
+        self.engine = create_engine('postgresql://postgres:postgres@localhost:6432/cipherstash_getting_started')
+        Session = sessionmaker(bind=self.engine)
+        self.session = Session()
+        BaseModel.metadata.create_all(self.engine)
+
+        self.session.query(Example).delete()
+        self.example = Example(
+            e_int=1, e_utf8_str="str", e_jsonb='{"key": "value"}', e_float=1.1, e_date=date(2024, 1, 1), e_bool=True
+        )
+        self.session.add(self.example)
+        self.session.commit()
+
+    def test_encrypted_int(self):
+        found = self.session.query(Example).filter(Example.id == self.example.id).one()
+        self.assertEqual(found.encrypted_int, 1)
+
+    def test_encrypted_boolean(self):
+        found = self.session.query(Example).filter(Example.id == self.example.id).one()
+        self.assertEqual(found.encrypted_boolean, True)
+
+    def test_encrypted_date(self):
+        found = self.session.query(Example).filter(Example.id == self.example.id).one()
+        self.assertEqual(found.encrypted_date, date(2024, 1, 1))
+
+    def test_encrypted_float(self):
+        found = self.session.query(Example).filter(Example.id == self.example.id).one()
+        self.assertEqual(found.encrypted_float, 1.1)
+    
+    def test_encrypted_utf8_str(self):
+        found = self.session.query(Example).filter(Example.id == self.example.id).one()
+        self.assertEqual(found.encrypted_utf8_str, "str")
+
+    def test_encrypted_jsonb(self):
+        found = self.session.query(Example).filter(Example.id == self.example.id).one()
+        self.assertEqual(found.encrypted_jsonb,  {"key": "value"})
+    
+if __name__ == '__main__':
+    unittest.main()

languages/python/jupyter_notebook/cs_types.py

@@ -1,6 +1,3 @@
-#!/usr/bin/env python
-
-import psycopg2
 from psycopg2.extras import RealDictCursor
 from pprint import pprint
 from datetime import datetime
@@ -12,71 +9,69 @@ def __init__(self, v, t: str, c: str):
         self.table = t
         self.column = c
 
-    def to_db_format(self):
+    def to_db_format(self, query_type=None):
         data = {
             "k": "pt",
-            "p": self.value_in_db_format(),
+            "p": self.value_in_db_format(query_type),
             "i": {
               "t": str(self.table),
               "c": str(self.column)
             },
             "v": 1,
+            "q": query_type,
         }
         return json.dumps(data)
 
-    # TODO: Unused at the moment
-    @classmethod
-    def from_json_str(cls, json_str):
-        parsed = json.loads(json_str)
-        return cls.from_parsed_json(parsed)
-
     @classmethod
     def from_parsed_json(cls, parsed):
         return cls.value_from_db_format(parsed["p"])
 
 class CsInt(CsValue):
-    def value_in_db_format(self):
+    def value_in_db_format(self, query_type = None):
         return str(self.value)
 
     @classmethod
     def value_from_db_format(cls, s: str):
         return int(s)
 
 class CsBool(CsValue):
-    def value_in_db_format(self):
+    def value_in_db_format(self, query_type = None):
         return str(self.value).lower()
 
     @classmethod
     def value_from_db_format(cls, s: str):
         return s.lower() == 'true'
 
 class CsDate(CsValue):
-    def value_in_db_format(self):
+    def value_in_db_format(self, query_type = None):
         return self.value.isoformat()
 
     @classmethod
     def value_from_db_format(cls, s: str):
-        return datetime.fromisoformat(s)
+        return datetime.fromisoformat(s).date()
 
 class CsFloat(CsValue):
-    def value_in_db_format(self):
+    def value_in_db_format(self, query_type = None):
         return str(self.value)
 
     @classmethod
     def value_from_db_format(cls, s: str):
         return float(s)
 
 class CsText(CsValue):
-    def value_in_db_format(self):
+    def value_in_db_format(self, query_type = None):
         return self.value
 
     @classmethod
     def value_from_db_format(cls, s: str):
         return s
 
 class CsJsonb(CsValue):
-    def value_in_db_format(self):
-        return json.dumps(self.value)
+    def value_in_db_format(self, query_type):
+        if query_type == "ejson_path":
+            return self.value
+        else:
+            return json.dumps(self.value)
 
     @classmethod
     def value_from_db_format(cls, s: str):

languages/python/jupyter_notebook/cs_types_test.py

@@ -0,0 +1,130 @@
+import unittest
+import json
+from datetime import date
+from cs_types import *
+
+class EqlTest(unittest.TestCase):
+    def test(self):
+        self.assertTrue(True)
+
+    def test_to_db_format(self):
+        self.assertEqual(
+            CsInt(1, "table", "column").to_db_format(),
+            '{"k": "pt", "p": "1", "i": {"t": "table", "c": "column"}, "v": 1, "q": null}'
+        )
+
+    def test_from_parsed_json_uses_p_value(self):
+        parsed = json.loads('{"k": "pt", "p": "1", "i": {"t": "table", "c": "column"}, "v": 1, "q": null}')
+        self.assertEqual(
+            CsInt.from_parsed_json(parsed),
+            1
+        )
+
+    def test_cs_int_prints_value(self):
+        cs_int = CsInt(1, "table", "column")
+        self.assertEqual(
+            cs_int.value_in_db_format(),
+            "1"
+        )
+
+    def test_ces_int_makes_int(self):
+        self.assertEqual(
+            CsInt.value_from_db_format("1"),
+            1
+        )
+
+    def test_cs_bool_prints_value_in_lower_case(self):
+        cs_bool = CsBool(True, "table", "column")
+        self.assertEqual(
+            cs_bool.value_in_db_format(),
+            "true"
+        )
+    
+    def test_cs_bool_returns_bool(self):
+        self.assertEqual(
+            CsBool.value_from_db_format("true"),
+            True
+        )
+
+    def test_cs_date_prints_value(self):
+        cs_date = CsDate(date(2024, 11, 1), "table", "column")
+        self.assertEqual(
+            cs_date.value_in_db_format(),
+            "2024-11-01"
+        )
+
+    def test_cs_date_returns_datetime(self):
+        self.assertEqual(
+            CsDate.value_from_db_format("2024-11-01"),
+            date(2024, 11, 1)
+        )
+
+    def test_cs_float_prints_value(self):
+        cs_float = CsFloat(1.1, "table", "column")
+        self.assertEqual(
+            cs_float.value_in_db_format(),
+            "1.1"
+        )
+
+    def test_cs_float_returns_float(self):
+        self.assertEqual(
+            CsFloat.value_from_db_format("1.1"),
+            1.1
+        )
+
+    def test_cs_text_prints_value(self):
+        cs_text = CsText("text", "table", "column")
+        self.assertEqual(
+            cs_text.value_in_db_format(),
+            "text"
+        )
+    
+    def test_cs_text_returns_value(self):
+        self.assertEqual(
+            CsText.value_from_db_format("text"),
+            "text"
+        )
+
+    def test_cs_jsonb_prints_json_string(self):
+        cs_jsonb = CsJsonb({"a": 1}, "table", "column")
+        self.assertEqual(
+            cs_jsonb.value_in_db_format("ste_vec"),
+            '{"a": 1}'
+        )
+
+    def test_cs_jsonb_prints_value_for_ejson_path(self):
+        cs_jsonb = CsJsonb("$.a.b", "table", "column")
+        self.assertEqual(
+            cs_jsonb.value_in_db_format("ejson_path"),
+            '$.a.b'
+        )
+
+    def test_cs_jsonb_returns_value(self):
+        self.assertEqual(
+            CsJsonb.value_from_db_format('{"a": 1}'),
+            {"a": 1}
+        )
+
+    def test_cs_row_makes_row(self):
+        cs_row = CsRow(
+            {"encrypted_int": json.loads(CsInt(1, "table", "column").to_db_format()),
+             "encrypted_boolean": json.loads(CsBool(True, "table", "column").to_db_format()),
+             "encrypted_date": json.loads(CsDate(date(2024, 11, 1), "table", "column").to_db_format()),
+             "encrypted_float": json.loads(CsFloat(1.1, "table", "column").to_db_format()),
+             "encrypted_utf8_str": json.loads(CsText("text", "table", "column").to_db_format()),
+             "encrypted_jsonb": json.loads(CsJsonb('{"a": 1}', "table", "column").to_db_format())
+            })
+
+        self.assertEqual(
+            cs_row.row,
+            {"encrypted_int": 1,
+             "encrypted_boolean": True,
+             "encrypted_date": date(2024, 11, 1),
+             "encrypted_float": 1.1,
+             "encrypted_utf8_str": "text",
+             "encrypted_jsonb": '"{\\"a\\": 1}"'
+            }
+        )
+
+if __name__ == '__main__':
+    unittest.main()