Add encryption_scheme property to LeapDocument.

4 files changed, 66 insertions, 3 deletions

diff --git a/src/leap/soledad/tests/__init__.py b/src/leap/soledad/tests/__init__.py
index 396b2775..28114391 100644
--- a/src/leap/soledad/tests/__init__.py
+++ b/src/leap/soledad/tests/__init__.py
@@ -3,6 +3,8 @@ Tests to make sure Soledad provides U1DB functionality and more.
 """
 
 import u1db
+
+
 from leap.soledad import Soledad
 from leap.soledad.crypto import SoledadCrypto
 from leap.soledad.backends.leap_backend import LeapDocument
diff --git a/src/leap/soledad/tests/test_couch.py b/src/leap/soledad/tests/test_couch.py
index 008c3ca4..cdf9c9ff 100644
--- a/src/leap/soledad/tests/test_couch.py
+++ b/src/leap/soledad/tests/test_couch.py
@@ -15,6 +15,10 @@ try:
     import simplejson as json
 except ImportError:
     import json  # noqa
+from leap.soledad.backends.leap_backend import (
+    LeapDocument,
+    EncryptionSchemes,
+)
 
 
 #-----------------------------------------------------------------------------
@@ -164,10 +168,14 @@ def copy_couch_database_for_test(test, db):
     return new_db
 
 
+def make_document_for_test(test, doc_id, rev, content, has_conflicts=False):
+    return LeapDocument(doc_id, rev, content, has_conflicts=has_conflicts)
+
+
 COUCH_SCENARIOS = [
     ('couch', {'make_database_for_test': make_couch_database_for_test,
                'copy_database_for_test': copy_couch_database_for_test,
-               'make_document_for_test': tests.make_document_for_test, }),
+               'make_document_for_test': make_document_for_test, }),
 ]
 
 
@@ -404,5 +412,30 @@ class CouchDatabaseStorageTests(CouchDBTestCase):
         content = self._fetch_u1db_data(db)
         self.assertEqual(db._replica_uid, content['replica_uid'])
 
+    def test_store_encryption_scheme(self):
+        db = couch.CouchDatabase('http://localhost:' + str(self.wrapper.port),
+                                 'u1db_tests')
+        doc = db.create_doc_from_json(tests.simple_doc)
+        # assert that docs have no encryption_scheme by default
+        self.assertEqual(EncryptionSchemes.NONE, doc.encryption_scheme)
+        # assert that we can store a different value
+        doc.encryption_scheme = EncryptionSchemes.PUBKEY
+        db.put_doc(doc)
+        self.assertEqual(
+            EncryptionSchemes.PUBKEY,
+            db.get_doc(doc.doc_id).encryption_scheme)
+        # assert that we can store another different value
+        doc.encryption_scheme = EncryptionSchemes.SYMKEY
+        db.put_doc(doc)
+        self.assertEqual(
+            EncryptionSchemes.SYMKEY,
+            db.get_doc(doc.doc_id).encryption_scheme)
+        # assert that we can store the default value
+        doc.encryption_scheme = EncryptionSchemes.NONE
+        db.put_doc(doc)
+        self.assertEqual(
+            EncryptionSchemes.NONE,
+            db.get_doc(doc.doc_id).encryption_scheme)
+
 
 load_tests = tests.load_with_scenarios
diff --git a/src/leap/soledad/tests/test_leap_backend.py b/src/leap/soledad/tests/test_leap_backend.py
index fd9ef85d..0bf70f9e 100644
--- a/src/leap/soledad/tests/test_leap_backend.py
+++ b/src/leap/soledad/tests/test_leap_backend.py
@@ -5,6 +5,8 @@ For these tests to run, a leap server has to be running on (default) port
 """
 
 import u1db
+
+
 from leap.soledad.backends import leap_backend
 from leap.soledad.tests import u1db_tests as tests
 from leap.soledad.tests.u1db_tests.test_remote_sync_target import (
diff --git a/src/leap/soledad/tests/test_sqlcipher.py b/src/leap/soledad/tests/test_sqlcipher.py
index b71478ac..337ab4ee 100644
--- a/src/leap/soledad/tests/test_sqlcipher.py
+++ b/src/leap/soledad/tests/test_sqlcipher.py
@@ -21,7 +21,10 @@ from leap.soledad.backends.sqlcipher import (
     DatabaseIsNotEncrypted,
 )
 from leap.soledad.backends.sqlcipher import open as u1db_open
-from leap.soledad.backends.leap_backend import LeapDocument
+from leap.soledad.backends.leap_backend import (
+    LeapDocument,
+    EncryptionSchemes,
+)
 
 # u1db tests stuff.
 from leap.soledad.tests import u1db_tests as tests
@@ -342,6 +345,29 @@ class TestSQLCipherPartialExpandDatabase(
         self.db.put_doc(doc)
         self.assertEqual(True, self.db.get_doc(doc.doc_id).syncable)
 
+    def test_store_encryption_scheme(self):
+        doc = self.db.create_doc_from_json(tests.simple_doc)
+        # assert that docs have no encryption_scheme by default
+        self.assertEqual(EncryptionSchemes.NONE, doc.encryption_scheme)
+        # assert that we can store a different value
+        doc.encryption_scheme = EncryptionSchemes.PUBKEY
+        self.db.put_doc(doc)
+        self.assertEqual(
+            EncryptionSchemes.PUBKEY,
+            self.db.get_doc(doc.doc_id).encryption_scheme)
+        # assert that we can store another different value
+        doc.encryption_scheme = EncryptionSchemes.SYMKEY
+        self.db.put_doc(doc)
+        self.assertEqual(
+            EncryptionSchemes.SYMKEY,
+            self.db.get_doc(doc.doc_id).encryption_scheme)
+        # assert that we can store the default value
+        doc.encryption_scheme = EncryptionSchemes.NONE
+        self.db.put_doc(doc)
+        self.assertEqual(
+            EncryptionSchemes.NONE,
+            self.db.get_doc(doc.doc_id).encryption_scheme)
+
 
 #-----------------------------------------------------------------------------
 # The following tests come from `u1db.tests.test_open`.
@@ -413,7 +439,7 @@ def sync_via_synchronizer_and_leap(test, db_source, db_target,
 sync_scenarios.append(('pyleap', {
     'make_database_for_test': test_sync.make_database_for_http_test,
     'copy_database_for_test': test_sync.copy_database_for_http_test,
-    'make_document_for_test': tests.make_document_for_test,
+    'make_document_for_test': make_document_for_test,
     'make_app_with_state': tests.test_remote_sync_target.make_http_app,
     'do_sync': sync_via_synchronizer_and_leap,
 }))