Repository: iterative/dvc

Issue:
state: store inode as int and not as TEXT
Kudos @shcheklein for noticing.


Discussion:

Resolving patch:
diff --git a/dvc/project.py b/dvc/project.py
--- a/dvc/project.py
+++ b/dvc/project.py
@@ -22,7 +22,7 @@ class Project(object):
     def __init__(self, root_dir):
         from dvc.logger import Logger
         from dvc.config import Config
-        from dvc.state import LinkState, State
+        from dvc.state import State
         from dvc.lock import Lock
         from dvc.scm import SCM
         from dvc.cache import Cache
@@ -39,7 +39,6 @@ def __init__(self, root_dir):
         # NOTE: storing state and link_state in the repository itself to avoid
         # any possible state corruption in 'shared cache dir' scenario.
         self.state = State(self)
-        self.link_state = LinkState(self)
 
         core = self.config._config[Config.SECTION_CORE]
         self.logger = Logger(core.get(Config.SECTION_CORE_LOGLEVEL, None))
@@ -149,16 +148,17 @@ def add(self, fname, recursive=False):
             fnames = [fname]
 
         stages = []
-        for f in fnames:
-            stage = Stage.loads(project=self,
-                                outs=[f],
-                                add=True)
+        with self.state:
+            for f in fnames:
+                stage = Stage.loads(project=self,
+                                    outs=[f],
+                                    add=True)
 
-            self._check_output_duplication(stage.outs)
+                self._check_output_duplication(stage.outs)
 
-            stage.save()
-            stage.dump()
-            stages.append(stage)
+                stage.save()
+                stage.dump()
+                stages.append(stage)
 
         return stages
 
@@ -185,34 +185,35 @@ def move(self, from_path, to_path):
                                      [from_path])[0]
 
         found = False
-        for stage in self.stages():
-            for out in stage.outs:
-                if out.path != from_out.path:
-                    continue
+        with self.state:
+            for stage in self.stages():
+                for out in stage.outs:
+                    if out.path != from_out.path:
+                        continue
 
-                if not stage.is_data_source:
-                    msg = 'Dvcfile \'{}\' is not a data source.'
-                    raise DvcException(msg.format(stage.rel_path))
+                    if not stage.is_data_source:
+                        msg = 'Dvcfile \'{}\' is not a data source.'
+                        raise DvcException(msg.format(stage.rel_path))
 
-                found = True
-                to_out = Output.loads_from(out.stage,
-                                           [to_path],
-                                           out.cache,
-                                           out.metric)[0]
-                out.move(to_out)
+                    found = True
+                    to_out = Output.loads_from(out.stage,
+                                               [to_path],
+                                               out.cache,
+                                               out.metric)[0]
+                    out.move(to_out)
 
-                stage_base = os.path.basename(stage.path)
-                stage_base = stage_base.rstrip(Stage.STAGE_FILE_SUFFIX)
+                    stage_base = os.path.basename(stage.path)
+                    stage_base = stage_base.rstrip(Stage.STAGE_FILE_SUFFIX)
 
-                stage_dir = os.path.dirname(stage.path)
-                from_base = os.path.basename(from_path)
-                to_base = os.path.basename(to_path)
-                if stage_base == from_base:
-                    os.unlink(stage.path)
-                    path = to_base + Stage.STAGE_FILE_SUFFIX
-                    stage.path = os.path.join(stage_dir, path)
+                    stage_dir = os.path.dirname(stage.path)
+                    from_base = os.path.basename(from_path)
+                    to_base = os.path.basename(to_path)
+                    if stage_base == from_base:
+                        os.unlink(stage.path)
+                        path = to_base + Stage.STAGE_FILE_SUFFIX
+                        stage.path = os.path.join(stage_dir, path)
 
-            stage.dump()
+                stage.dump()
 
         if not found:
             msg = 'Unable to find dvcfile with output \'{}\''
@@ -240,9 +241,12 @@ def run(self,
 
         self._check_output_duplication(stage.outs)
 
-        if not no_exec:
-            stage.run()
+        with self.state:
+            if not no_exec:
+                stage.run()
+
         stage.dump()
+
         return stage
 
     def imp(self, url, out):
@@ -253,7 +257,9 @@ def imp(self, url, out):
 
         self._check_output_duplication(stage.outs)
 
-        stage.run()
+        with self.state:
+            stage.run()
+
         stage.dump()
         return stage
 
@@ -292,15 +298,22 @@ def reproduce(self,
             core = config._config[config.SECTION_CORE]
             interactive = core.get(config.SECTION_CORE_INTERACTIVE, False)
 
-        if recursive:
-            return self._reproduce_stages(G,
-                                          stages,
-                                          node,
-                                          force,
-                                          dry,
-                                          interactive)
+        with self.state:
+            if recursive:
+                ret = self._reproduce_stages(G,
+                                             stages,
+                                             node,
+                                             force,
+                                             dry,
+                                             interactive)
+            else:
+                ret = self._reproduce_stage(stages,
+                                            node,
+                                            force,
+                                            dry,
+                                            interactive)
 
-        return self._reproduce_stage(stages, node, force, dry, interactive)
+        return ret
 
     def _reproduce_stages(self, G, stages, node, force, dry, interactive):
         import networkx as nx
@@ -322,7 +335,7 @@ def _cleanup_unused_links(self, all_stages):
         for stage in all_stages:
             for out in stage.outs:
                 used.append(out.path)
-        self.link_state.remove_unused(used)
+        self.state.remove_unused_links(used)
 
     def checkout(self, target=None):
         all_stages = self.active_stages()
@@ -332,15 +345,16 @@ def checkout(self, target=None):
         else:
             stages = all_stages
 
-        self._cleanup_unused_links(all_stages)
+        with self.state:
+            self._cleanup_unused_links(all_stages)
 
-        for stage in stages:
-            if stage.locked:
-                msg = 'DVC file \'{}\' is locked. Its dependecies are not ' \
-                      'going to be checked out.'
-                self.logger.warn(msg.format(stage.relpath))
+            for stage in stages:
+                if stage.locked:
+                    msg = 'DVC file \'{}\' is locked. Its dependecies are ' \
+                          'not going to be checked out.'
+                    self.logger.warn(msg.format(stage.relpath))
 
-            stage.checkout()
+                stage.checkout()
 
     def _used_cache(self, target=None, all_branches=False, active=True):
         cache = {}
@@ -381,29 +395,30 @@ def _do_gc(self, typ, func, clist):
             self.logger.info("No unused {} cache to remove.".format(typ))
 
     def gc(self, all_branches=False, cloud=False, remote=None):
-        clist = self._used_cache(target=None,
-                                 all_branches=all_branches,
-                                 active=False)
-        self._do_gc('local', self.cache.local.gc, clist)
+        with self.state:
+            clist = self._used_cache(target=None,
+                                     all_branches=all_branches,
+                                     active=False)
+            self._do_gc('local', self.cache.local.gc, clist)
 
-        if self.cache.s3:
-            self._do_gc('s3', self.cache.s3.gc, clist)
+            if self.cache.s3:
+                self._do_gc('s3', self.cache.s3.gc, clist)
 
-        if self.cache.gs:
-            self._do_gc('gs', self.cache.gs.gc, clist)
+            if self.cache.gs:
+                self._do_gc('gs', self.cache.gs.gc, clist)
 
-        if self.cache.ssh:
-            self._do_gc('ssh', self.cache.ssh.gc, clist)
+            if self.cache.ssh:
+                self._do_gc('ssh', self.cache.ssh.gc, clist)
 
-        if self.cache.hdfs:
-            self._do_gc('hdfs', self.cache.hdfs.gc, clist)
+            if self.cache.hdfs:
+                self._do_gc('hdfs', self.cache.hdfs.gc, clist)
 
-        if self.cache.azure:
-            self._do_gc('azure', self.cache.azure.gc, clist)
+            if self.cache.azure:
+                self._do_gc('azure', self.cache.azure.gc, clist)
 
-        if cloud:
-            self._do_gc('remote', self.cloud._get_cloud(remote,
-                                                        'gc -c').gc, clist)
+            if cloud:
+                self._do_gc('remote', self.cloud._get_cloud(remote,
+                                                            'gc -c').gc, clist)
 
     def push(self,
              target=None,
@@ -411,10 +426,11 @@ def push(self,
              remote=None,
              all_branches=False,
              show_checksums=False):
-        self.cloud.push(self._used_cache(target, all_branches)['local'],
-                        jobs,
-                        remote=remote,
-                        show_checksums=show_checksums)
+        with self.state:
+            self.cloud.push(self._used_cache(target, all_branches)['local'],
+                            jobs,
+                            remote=remote,
+                            show_checksums=show_checksums)
 
     def fetch(self,
               target=None,
@@ -422,10 +438,11 @@ def fetch(self,
               remote=None,
               all_branches=False,
               show_checksums=False):
-        self.cloud.pull(self._used_cache(target, all_branches)['local'],
-                        jobs,
-                        remote=remote,
-                        show_checksums=show_checksums)
+        with self.state:
+            self.cloud.pull(self._used_cache(target, all_branches)['local'],
+                            jobs,
+                            remote=remote,
+                            show_checksums=show_checksums)
 
     def pull(self,
              target=None,
@@ -488,12 +505,13 @@ def status(self,
                cloud=False,
                remote=None,
                show_checksums=False):
-        if cloud:
-            return self._cloud_status(target,
-                                      jobs,
-                                      remote=remote,
-                                      show_checksums=show_checksums)
-        return self._local_status(target)
+        with self.state:
+            if cloud:
+                return self._cloud_status(target,
+                                          jobs,
+                                          remote=remote,
+                                          show_checksums=show_checksums)
+            return self._local_status(target)
 
     def _read_metric_json(self, fd, json_path):
         import json
diff --git a/dvc/remote/local.py b/dvc/remote/local.py
--- a/dvc/remote/local.py
+++ b/dvc/remote/local.py
@@ -36,7 +36,6 @@ class RemoteLOCAL(RemoteBase):
     def __init__(self, project, config):
         self.project = project
         self.state = self.project.state
-        self.link_state = project.link_state
         storagepath = config.get(Config.SECTION_AWS_STORAGEPATH, None)
         self.cache_dir = config.get(Config.SECTION_REMOTE_URL, storagepath)
 
@@ -138,7 +137,6 @@ def ntpath(path):
     def collect_dir_cache(self, dname):
         dir_info = []
 
-        db = self.state.load()
         bar = False
         for root, dirs, files in os.walk(dname):
             if len(files) > LARGE_DIR_SIZE:
@@ -159,13 +157,10 @@ def collect_dir_cache(self, dname):
                     progress.update_target(title, processed, total)
                     processed += 1
 
-                md5 = self.state.update(path, use_db=db)
+                md5 = self.state.update(path)
                 dir_info.append({self.PARAM_RELPATH: relpath,
                                  self.PARAM_MD5: md5})
 
-        db.commit()
-        db.close()
-
         if bar:
             progress.finish_target(title)
 
@@ -253,7 +248,7 @@ def checkout(self, path_info, checksum_info):
 
         if not self.is_dir_cache(cache):
             self.link(cache, path)
-            self.link_state.update(path)
+            self.state.update_link(path)
             return
 
         # Create dir separately so that dir is created
@@ -267,7 +262,7 @@ def checkout(self, path_info, checksum_info):
             relpath = entry[self.PARAM_RELPATH]
             p = os.path.join(path, relpath)
             self.link(c, p)
-        self.link_state.update(path)
+        self.state.update_link(path)
 
     def _move(self, inp, outp):
         # moving in two stages to make last the move atomic in
@@ -289,7 +284,7 @@ def _save_file(self, path_info):
             remove(path)
 
         self.link(cache, path)
-        self.link_state.update(path)
+        self.state.update_link(path)
 
         return {self.PARAM_MD5: md5}
 
@@ -310,7 +305,7 @@ def _save_dir(self, path_info):
 
             self.link(c, p)
 
-        self.link_state.update(path)
+        self.state.update_link(path)
 
         return {self.PARAM_MD5: md5}
 
diff --git a/dvc/state.py b/dvc/state.py
--- a/dvc/state.py
+++ b/dvc/state.py
@@ -1,4 +1,5 @@
 import os
+import time
 import sqlite3
 import nanotime
 
@@ -15,28 +16,34 @@ class StateDuplicateError(DvcException):
 class State(object):
     STATE_FILE = 'state'
     STATE_TABLE = 'state'
-    STATE_TABLE_LAYOUT = "inode TEXT PRIMARY KEY, " \
+    STATE_TABLE_LAYOUT = "inode INTEGER PRIMARY KEY, " \
                          "mtime TEXT NOT NULL, " \
-                         "md5 TEXT NOT NULL"
+                         "md5 TEXT NOT NULL, " \
+                         "timestamp TEXT NOT NULL"
+
+    LINK_STATE_TABLE = 'link_state'
+    LINK_STATE_TABLE_LAYOUT = "path TEXT PRIMARY KEY, " \
+                              "inode INTEGER NOT NULL, " \
+                              "mtime TEXT NOT NULL"
 
     def __init__(self, project):
         self.project = project
         self.dvc_dir = project.dvc_dir
+        self.root_dir = project.root_dir
 
         if not self.dvc_dir:
             self.state_file = None
             return
 
         self.state_file = os.path.join(self.dvc_dir, self.STATE_FILE)
-        # Try loading once to check that the file is indeed a database
-        # and reformat it if it is not.
-        try:
-            db = self.load()
-            db.close()
-        except sqlite3.DatabaseError:
-            os.unlink(self.state_file)
-            db = self.load()
-            db.close()
+        self.db = None
+        self.c = None
+
+    def __enter__(self):
+        self.load()
+
+    def __exit__(self, type, value, tb):
+        self.dump()
 
     @staticmethod
     def init(project):
@@ -57,15 +64,40 @@ def changed(self, path, md5):
         return actual.split('.')[0] != md5.split('.')[0]
 
     def load(self):
-        db = sqlite3.connect(self.state_file)
-        c = db.cursor()
-        # Check that the state file is indeed a database
-        cmd = "CREATE TABLE IF NOT EXISTS {} ({})"
-        c.execute(cmd.format(self.STATE_TABLE, self.STATE_TABLE_LAYOUT))
-        c.execute("SELECT name FROM sqlite_master WHERE type='table';")
-        db.commit()
-        c.close()
-        return db
+        retries = 1
+        while True:
+            assert self.db is None
+            self.db = sqlite3.connect(self.state_file)
+            self.c = self.db.cursor()
+
+            # Try loading once to check that the file is indeed a database
+            # and reformat it if it is not.
+            try:
+                # Check that the state file is indeed a database
+                cmd = "CREATE TABLE IF NOT EXISTS {} ({})"
+                self.c.execute(cmd.format(self.STATE_TABLE,
+                                          self.STATE_TABLE_LAYOUT))
+                self.c.execute(cmd.format(self.LINK_STATE_TABLE,
+                                          self.LINK_STATE_TABLE_LAYOUT))
+                return
+            except sqlite3.DatabaseError:
+                self.db.close()
+                self.c.close()
+                self.db = None
+                self.c = None
+                if retries > 0:
+                    os.unlink(self.state_file)
+                    retries -= 1
+                else:
+                    raise
+
+    def dump(self):
+        assert self.db is not None
+        self.db.commit()
+        self.c.close()
+        self.db.close()
+        self.db = None
+        self.c = None
 
     @staticmethod
     def mtime(path):
@@ -82,99 +114,63 @@ def mtime(path):
 
     @staticmethod
     def inode(path):
-        return str(System.inode(path))
+        return System.inode(path)
 
-    def _do_update(self, path, use_db=None):
+    def _do_update(self, path):
         if not os.path.exists(path):
             return (None, None)
 
         mtime = self.mtime(path)
         inode = self.inode(path)
 
-        if use_db is None:
-            db = self.load()
-        else:
-            db = use_db
-
-        c = db.cursor()
-
-        md5 = self._get(inode, mtime, use_db=db)
-        if md5:
-            return (md5, None)
-
-        md5, info = self._collect(path)
-
-        cmd = 'REPLACE INTO {}(inode, mtime, md5) ' \
-              'VALUES ("{}", "{}", "{}")'.format(self.STATE_TABLE,
-                                                 inode,
-                                                 mtime,
-                                                 md5)
+        cmd = 'SELECT * from {} WHERE inode={}'.format(self.STATE_TABLE,
+                                                       inode)
 
-        c.execute(cmd)
-        c.close()
-        if use_db is None:
-            db.commit()
-            db.close()
+        self.c.execute(cmd)
+        ret = self.c.fetchall()
+        if len(ret) == 0:
+            md5, info = self._collect(path)
+            cmd = 'INSERT INTO {}(inode, mtime, md5, timestamp) ' \
+                  'VALUES ({}, "{}", "{}", "{}")'
+            self.c.execute(cmd.format(self.STATE_TABLE,
+                                      inode,
+                                      mtime,
+                                      md5,
+                                      int(nanotime.timestamp(time.time()))))
+        else:
+            assert len(ret) == 1
+            assert len(ret[0]) == 4
+            i, m, md5, timestamp = ret[0]
+            assert i == inode
+            if mtime != m:
+                md5, info = self._collect(path)
+                cmd = 'UPDATE {} SET ' \
+                      'mtime = "{}", md5 = "{}", timestamp = "{}" ' \
+                      'WHERE inode = {}'
+                self.c.execute(cmd.format(self.STATE_TABLE,
+                                          mtime,
+                                          md5,
+                                          int(nanotime.timestamp(time.time())),
+                                          inode))
+            else:
+                info = None
+                cmd = 'UPDATE {} SET timestamp = "{}" WHERE inode = {}'
+                self.c.execute(cmd.format(self.STATE_TABLE,
+                                          int(nanotime.timestamp(time.time())),
+                                          inode))
 
         return (md5, info)
 
-    def update(self, path, use_db=None):
-        return self._do_update(path, use_db=use_db)[0]
+    def update(self, path):
+        return self._do_update(path)[0]
 
-    def update_info(self, path, use_db=None):
-        md5, info = self._do_update(path, use_db=use_db)
+    def update_info(self, path):
+        md5, info = self._do_update(path)
         if not info:
             info = self.project.cache.local.load_dir_cache(md5)
         return (md5, info)
 
-    def _get(self, inode, mtime, use_db=None):
-        cmd = 'SELECT * from {} WHERE inode="{}"'.format(self.STATE_TABLE,
-                                                         inode)
-
-        if use_db is None:
-            db = self.load()
-        else:
-            db = use_db
-
-        c = db.cursor()
-        c.execute(cmd)
-        ret = c.fetchall()
-        c.close()
-
-        if use_db is None:
-            db.commit()
-            db.close()
-
-        if len(ret) == 0:
-            return None
-        assert len(ret) == 1
-        assert len(ret[0]) == 3
-        i, m, md5 = ret[0]
-        assert i == inode
-
-        if mtime == m:
-            return md5
-
-        return None
-
-    def get(self, path):
-        mtime = self.mtime(path)
-        inode = self.inode(path)
-
-        return self._get(inode, mtime)
-
-
-class LinkState(State):
-    STATE_TABLE = 'link'
-    STATE_TABLE_LAYOUT = "path TEXT PRIMARY KEY, " \
-                         "inode TEXT NOT NULL, " \
-                         "mtime TEXT NOT NULL"
-
-    def __init__(self, project):
-        super(LinkState, self).__init__(project)
-        self.root_dir = project.root_dir
-
-    def update(self, path, use_db=None):
+    def update_link(self, path):
         if not os.path.exists(path):
             return
 
@@ -182,32 +178,18 @@ def update(self, path, use_db=None):
         inode = self.inode(path)
         relpath = os.path.relpath(path, self.root_dir)
 
-        if use_db is None:
-            db = self.load()
-        else:
-            db = use_db
-
-        c = db.cursor()
-
         cmd = 'REPLACE INTO {}(path, inode, mtime) ' \
-              'VALUES ("{}", "{}", "{}")'.format(self.STATE_TABLE,
-                                                 relpath,
-                                                 inode,
-                                                 mtime)
-        c.execute(cmd)
-        c.close()
-
-        if use_db is None:
-            db.commit()
-            db.close()
-
-    def remove_unused(self, used):
+              'VALUES ("{}", {}, "{}")'.format(self.LINK_STATE_TABLE,
+                                               relpath,
+                                               inode,
+                                               mtime)
+        self.c.execute(cmd)
+
+    def remove_unused_links(self, used):
         unused = []
 
-        db = self.load()
-        c = db.cursor()
-        c.execute('SELECT * FROM {}'.format(self.STATE_TABLE))
-        for row in c:
+        self.c.execute('SELECT * FROM {}'.format(self.LINK_STATE_TABLE))
+        for row in self.c:
             p, i, m = row
             path = os.path.join(self.root_dir, p)
 
@@ -225,11 +207,6 @@ def remove_unused(self, used):
                 remove(path)
                 unused.append(p)
 
-        db.commit()
         for p in unused:
             cmd = 'DELETE FROM {} WHERE path = "{}"'
-            c.execute(cmd.format(self.STATE_TABLE, p))
-
-        c.close()
-        db.commit()
-        db.close()
+            self.c.execute(cmd.format(self.LINK_STATE_TABLE, p))


Test patch:
diff --git a/tests/test_data_cloud.py b/tests/test_data_cloud.py
--- a/tests/test_data_cloud.py
+++ b/tests/test_data_cloud.py
@@ -221,55 +221,56 @@ def _test_cloud(self):
         info_dir = stage_dir.outs[0].dumpd()
         md5_dir = info_dir['md5']
 
-        # Check status
-        sleep()
-        status = self.cloud.status([info], show_checksums=True)
-        self.assertEqual([(md5, STATUS_NEW)], status)
+        with self.cloud.project.state:
+            # Check status
+            sleep()
+            status = self.cloud.status([info], show_checksums=True)
+            self.assertEqual([(md5, STATUS_NEW)], status)
 
-        status_dir = self.cloud.status([info_dir], show_checksums=True)
-        self.assertTrue((md5_dir, STATUS_NEW) in status_dir)
+            status_dir = self.cloud.status([info_dir], show_checksums=True)
+            self.assertTrue((md5_dir, STATUS_NEW) in status_dir)
 
-        # Push and check status
-        self.cloud.push([info])
-        self.assertTrue(os.path.exists(cache))
-        self.assertTrue(os.path.isfile(cache))
+            # Push and check status
+            self.cloud.push([info])
+            self.assertTrue(os.path.exists(cache))
+            self.assertTrue(os.path.isfile(cache))
 
-        self.cloud.push([info_dir])
-        self.assertTrue(os.path.isfile(cache_dir))
+            self.cloud.push([info_dir])
+            self.assertTrue(os.path.isfile(cache_dir))
 
-        sleep()
-        status = self.cloud.status([info], show_checksums=True)
-        self.assertEqual([(md5, STATUS_OK)], status)
+            sleep()
+            status = self.cloud.status([info], show_checksums=True)
+            self.assertEqual([(md5, STATUS_OK)], status)
 
-        status_dir = self.cloud.status([info_dir], show_checksums=True)
-        self.assertTrue((md5_dir, STATUS_OK) in status_dir)
+            status_dir = self.cloud.status([info_dir], show_checksums=True)
+            self.assertTrue((md5_dir, STATUS_OK) in status_dir)
 
-        # Remove and check status
-        sleep()
-        shutil.rmtree(self.dvc.cache.local.cache_dir)
+            # Remove and check status
+            sleep()
+            shutil.rmtree(self.dvc.cache.local.cache_dir)
 
-        status = self.cloud.status([info], show_checksums=True)
-        self.assertEqual([(md5, STATUS_DELETED)], status)
+            status = self.cloud.status([info], show_checksums=True)
+            self.assertEqual([(md5, STATUS_DELETED)], status)
 
-        status_dir = self.cloud.status([info_dir], show_checksums=True)
-        self.assertEqual([(md5_dir, STATUS_DELETED)], status_dir)
+            status_dir = self.cloud.status([info_dir], show_checksums=True)
+            self.assertEqual([(md5_dir, STATUS_DELETED)], status_dir)
 
-        # Pull and check status
-        self.cloud.pull([info])
-        self.assertTrue(os.path.exists(cache))
-        self.assertTrue(os.path.isfile(cache))
-        with open(cache, 'r') as fd:
-            self.assertEqual(fd.read(), self.FOO_CONTENTS)
+            # Pull and check status
+            self.cloud.pull([info])
+            self.assertTrue(os.path.exists(cache))
+            self.assertTrue(os.path.isfile(cache))
+            with open(cache, 'r') as fd:
+                self.assertEqual(fd.read(), self.FOO_CONTENTS)
 
-        self.cloud.pull([info_dir])
-        self.assertTrue(os.path.isfile(cache_dir))
+            self.cloud.pull([info_dir])
+            self.assertTrue(os.path.isfile(cache_dir))
 
-        sleep()
-        status = self.cloud.status([info], show_checksums=True)
-        self.assertEqual([(md5, STATUS_OK)], status)
+            sleep()
+            status = self.cloud.status([info], show_checksums=True)
+            self.assertEqual([(md5, STATUS_OK)], status)
 
-        status_dir = self.cloud.status([info_dir], show_checksums=True)
-        self.assertTrue((md5_dir, STATUS_OK) in status_dir)
+            status_dir = self.cloud.status([info_dir], show_checksums=True)
+            self.assertTrue((md5_dir, STATUS_OK) in status_dir)
 
     def test(self):
         self._ensure_should_run()
@@ -335,9 +336,6 @@ def _get_url(self):
     def _get_cloud_class(self):
         return RemoteSSH
 
-    def test(self):
-        super(TestRemoteSSH, self).test()
-
 
 class TestRemoteHDFS(TestDataCloudBase):
     def _should_test(self):
@@ -349,9 +347,6 @@ def _get_url(self):
     def _get_cloud_class(self):
         return RemoteHDFS
 
-    def test(self):
-        super(TestRemoteHDFS, self).test()
-
 
 class TestDataCloudCLIBase(TestDvc):
     def main(self, args):
diff --git a/tests/test_state.py b/tests/test_state.py
--- a/tests/test_state.py
+++ b/tests/test_state.py
@@ -18,23 +18,24 @@ def test_update(self):
 
         state = State(self.dvc)
 
-        entry_md5 = state.update(path)
-        self.assertEqual(entry_md5, md5)
-
-        # Sleep some time to simulate realistic behavior.
-        # Some filesystems have a bad date resolution for
-        # mtime(i.e. 1sec for HFS) that cause problems with
-        # our 'state' system not being able to distinguish
-        # files that were modified within that delta.
-        time.sleep(1)        
-
-        os.unlink(path)
-        with open(path, 'w+') as fd:
-            fd.write('1')
-
-        md5 = file_md5(path)[0]
-        mtime = os.path.getmtime(path)
-        inode = System.inode(path)
-
-        entry_md5 = state.update(path)
-        self.assertEqual(entry_md5, md5)
+        with state:
+            entry_md5 = state.update(path)
+            self.assertEqual(entry_md5, md5)
+
+            # Sleep some time to simulate realistic behavior.
+            # Some filesystems have a bad date resolution for
+            # mtime(i.e. 1sec for HFS) that cause problems with
+            # our 'state' system not being able to distinguish
+            # files that were modified within that delta.
+            time.sleep(1)        
+
+            os.unlink(path)
+            with open(path, 'w+') as fd:
+                fd.write('1')
+
+            md5 = file_md5(path)[0]
+            mtime = os.path.getmtime(path)
+            inode = System.inode(path)
+
+            entry_md5 = state.update(path)
+            self.assertEqual(entry_md5, md5)