Snapshot specific symbol versions - Tests

Author: mstampfer

arctic/store/version_store.py

@@ -678,6 +678,8 @@ def snapshot(self, snap_name, metadata=None, skip_symbols=None, versions=None):
             an optional dictionary of metadata to persist along with the symbol.
         skip_symbols : `collections.Iterable`
             optional symbols to be excluded from the snapshot
+        versions: `dict`
+            an optional dictionary of versions of the symbols to be snapshot
         """
         # Ensure the user doesn't insert duplicates
         snapshot = self._snapshots.find_one({'name': snap_name})

tests/integration/store/test_version_store.py

@@ -491,6 +491,46 @@ def test_snapshot(library):
     assert versions[2]['snapshots'] == ['current']
 
 
+def test_snapshot_with_versions(library):
+    """ Test snapshot of write versions consistency """
+    library.write(symbol, ts1)
+    library.write(symbol, ts2)
+
+    # ensure snapshot of previous version is taken
+    library.snapshot('previous', versions={symbol: 1})
+    versions = library.list_versions(symbol)
+    assert versions[0]['snapshots'] == []
+    assert versions[1]['snapshots'] == ['previous']
+    assert_frame_equal(library.read(symbol, as_of='previous').data, ts1)
+
+    # ensure new snapshots are ordered after previous ones
+    library.snapshot('new')
+    versions = library.list_versions(symbol)
+    assert versions[0]['snapshots'] == ['new']
+    assert versions[0]['version'] == 2
+    assert_frame_equal(library.read(symbol, as_of='new').data, ts2)
+
+    assert versions[1]['snapshots'] == ['previous']
+    assert versions[1]['version'] == 1
+    assert_frame_equal(library.read(symbol, as_of='previous').data, ts1)
+
+    # ensure snapshot of previous version doesn't overwrite current version
+    library.write(symbol, ts1, prune_previous_version=True)
+    library.snapshot('another', versions={symbol: 1})
+    versions = library.list_versions(symbol)
+
+    assert versions[0]['snapshots'] == []
+    assert versions[0]['version'] == 3
+    assert_frame_equal(library.read(symbol).data, ts1)
+
+    assert versions[1]['snapshots'] == ['new']
+    assert versions[1]['version'] == 2
+
+    assert versions[2]['snapshots'] == ['previous', 'another']
+    assert versions[2]['version'] == 1
+    assert_frame_equal(library.read(symbol, as_of='another').data, ts1)
+
+
 def test_snapshot_exclusion(library):
     library.write(symbol, ts1)
     library.snapshot('current', skip_symbols=[symbol])

tests/unit/store/test_version_store.py

@@ -11,7 +11,7 @@
 from arctic.store import version_store
 from arctic.store.version_store import VersionStore, VersionedItem
 from arctic.arctic import ArcticLibraryBinding, Arctic
-from arctic.exceptions import ConcurrentModificationException
+from arctic.exceptions import ConcurrentModificationException, DuplicateSnapshotException
 from pymongo.errors import OperationFailure
 from pymongo.collection import Collection
 
@@ -202,3 +202,24 @@ def test_read_reports_random_errors():
             VersionStore.read(self, sentinel.symbol, sentinel.as_of, sentinel.from_version)
     assert 'bad' in str(e)
     assert le.call_count == 1
+
+
+def test_snapshot():
+    vs = create_autospec(VersionStore, _snapshots=Mock(),
+                                       _collection=Mock(),
+                                       _versions=Mock())
+    vs._snapshots.find_one.return_value = False
+    vs._versions.update_one.__name__ = 'name'
+    vs._snapshots.insert_one.__name__ = 'name'
+    vs.list_symbols.return_value = ['foo', 'bar']
+    VersionStore.snapshot(vs, "symbol")
+    assert vs._read_metadata.call_args_list == [call('foo', as_of=None, read_preference=ReadPreference.PRIMARY),
+                                                call('bar', as_of=None, read_preference=ReadPreference.PRIMARY)]
+
+
+def test_snapshot_duplicate_raises_exception():
+    vs = create_autospec(VersionStore, _snapshots=Mock())
+    with pytest.raises(DuplicateSnapshotException) as e:
+        vs._snapshots.find_one.return_value = True
+        VersionStore.snapshot(vs, 'symbol')
+        assert "Snapshot 'symbol' already exists" in str(e.value)