Merge "Raise VersionNotFoundError instead of KeyError"

3 files changed, 22 insertions, 3 deletions

diff --git a/migrate/exceptions.py b/migrate/exceptions.py
index cb8c409..31c8cd9 100644
--- a/migrate/exceptions.py
+++ b/migrate/exceptions.py
@@ -27,6 +27,10 @@ class InvalidVersionError(ControlledSchemaError):
     """Invalid version number."""
 
 
+class VersionNotFoundError(KeyError):
+    """Specified version is not present."""
+
+
 class DatabaseNotControlledError(ControlledSchemaError):
     """Database should be under version control, but it's not."""
 
diff --git a/migrate/tests/versioning/test_version.py b/migrate/tests/versioning/test_version.py
index df50072..286dd59 100644
--- a/migrate/tests/versioning/test_version.py
+++ b/migrate/tests/versioning/test_version.py
@@ -103,6 +103,10 @@ class TestVersion(fixture.Pathed):
         self.assertEqual(coll.latest, 4)
         self.assertEqual(len(coll.versions), 4)
         self.assertEqual(coll.version(4), coll.version(coll.latest))
+        # Check for non-existing version
+        self.assertRaises(VersionNotFoundError, coll.version, 5)
+        # Check for the current version
+        self.assertEqual('4', coll.version(4).version)
 
         coll2 = Collection(self.temp_usable_dir)
         self.assertEqual(coll.versions, coll2.versions)
diff --git a/migrate/versioning/version.py b/migrate/versioning/version.py
index 3ab814c..0633e1b 100644
--- a/migrate/versioning/version.py
+++ b/migrate/versioning/version.py
@@ -156,11 +156,22 @@ class Collection(pathed.Pathed):
             self.versions[ver].add_script(filepath)
 
     def version(self, vernum=None):
-        """Returns latest Version if vernum is not given.
-        Otherwise, returns wanted version"""
+        """Returns required version.
+
+        If vernum is not given latest version will be returned otherwise
+        required version will be returned.
+        :raises: : exceptions.VersionNotFoundError if respective migration
+        script file of version is not present in the migration repository.
+        """
         if vernum is None:
             vernum = self.latest
-        return self.versions[VerNum(vernum)]
+
+        try:
+            return self.versions[VerNum(vernum)]
+        except KeyError:
+            raise exceptions.VersionNotFoundError(
+                ("Database schema file with version %(args)s doesn't "
+                 "exist.") % {'args': VerNum(vernum)})
 
     @classmethod
     def clear(cls):