Repository: genialis/resolwe-bio-py

Issue:
Queries in new objects are not evaluated correctly
When new object is created, all queries that are included (i.e. `collection.data`, `collection.samples`,...) are constructed in a wrong way, because object's id is not yet known and filters are applied with `id=None`.

Queries shouldn't be initialized in `__init__`, but on the first call and error should be raised if `id` is not yet known.

Discussion:
codecov-io: # [Codecov](https://codecov.io/gh/genialis/resolwe-bio-py/pull/91?src=pr&el=h1) Report
> Merging [#91](https://codecov.io/gh/genialis/resolwe-bio-py/pull/91?src=pr&el=desc) into [master](https://codecov.io/gh/genialis/resolwe-bio-py/commit/afe16da59e19507196a3591650a71f94394ffe2e?src=pr&el=desc) will **increase** coverage by `0.03%`.
> The diff coverage is `86.45%`.


```diff
@@            Coverage Diff             @@
##           master      #91      +/-   ##
==========================================
+ Coverage    81.6%   81.64%   +0.03%     
==========================================
  Files          28       28              
  Lines        2153     2217      +64     
==========================================
+ Hits         1757     1810      +53     
- Misses        396      407      +11
```


| [Impacted Files](https://codecov.io/gh/genialis/resolwe-bio-py/pull/91?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [resdk/resolwe.py](https://codecov.io/gh/genialis/resolwe-bio-py/compare/afe16da59e19507196a3591650a71f94394ffe2e...194f027f7b6bacf730c5b5b2fa4bf6795f06bdde?src=pr&el=tree#diff-cmVzZGsvcmVzb2x3ZS5weQ==) | `98.88% <ø> (-0.01%)` | :x: |
| [resdk/scripts.py](https://codecov.io/gh/genialis/resolwe-bio-py/compare/afe16da59e19507196a3591650a71f94394ffe2e...194f027f7b6bacf730c5b5b2fa4bf6795f06bdde?src=pr&el=tree#diff-cmVzZGsvc2NyaXB0cy5weQ==) | `0% <ø> (ø)` | :white_check_mark: |
| [resdk/tests/unit/test_resolwe.py](https://codecov.io/gh/genialis/resolwe-bio-py/compare/afe16da59e19507196a3591650a71f94394ffe2e...194f027f7b6bacf730c5b5b2fa4bf6795f06bdde?src=pr&el=tree#diff-cmVzZGsvdGVzdHMvdW5pdC90ZXN0X3Jlc29sd2UucHk=) | `99.74% <100%> (ø)` | :white_check_mark: |
| [resdk/resources/sample.py](https://codecov.io/gh/genialis/resolwe-bio-py/compare/afe16da59e19507196a3591650a71f94394ffe2e...194f027f7b6bacf730c5b5b2fa4bf6795f06bdde?src=pr&el=tree#diff-cmVzZGsvcmVzb3VyY2VzL3NhbXBsZS5weQ==) | `84.21% <100%> (+4.79%)` | :white_check_mark: |
| [resdk/resources/data.py](https://codecov.io/gh/genialis/resolwe-bio-py/compare/afe16da59e19507196a3591650a71f94394ffe2e...194f027f7b6bacf730c5b5b2fa4bf6795f06bdde?src=pr&el=tree#diff-cmVzZGsvcmVzb3VyY2VzL2RhdGEucHk=) | `100% <100%> (ø)` | :white_check_mark: |
| [resdk/tests/unit/test_collections.py](https://codecov.io/gh/genialis/resolwe-bio-py/compare/afe16da59e19507196a3591650a71f94394ffe2e...194f027f7b6bacf730c5b5b2fa4bf6795f06bdde?src=pr&el=tree#diff-cmVzZGsvdGVzdHMvdW5pdC90ZXN0X2NvbGxlY3Rpb25zLnB5) | `99.01% <100%> (+0.16%)` | :white_check_mark: |
| [resdk/tests/unit/test_data.py](https://codecov.io/gh/genialis/resolwe-bio-py/compare/afe16da59e19507196a3591650a71f94394ffe2e...194f027f7b6bacf730c5b5b2fa4bf6795f06bdde?src=pr&el=tree#diff-cmVzZGsvdGVzdHMvdW5pdC90ZXN0X2RhdGEucHk=) | `98.92% <100%> (-0.06%)` | :x: |
| [resdk/tests/unit/test_query.py](https://codecov.io/gh/genialis/resolwe-bio-py/compare/afe16da59e19507196a3591650a71f94394ffe2e...194f027f7b6bacf730c5b5b2fa4bf6795f06bdde?src=pr&el=tree#diff-cmVzZGsvdGVzdHMvdW5pdC90ZXN0X3F1ZXJ5LnB5) | `99.4% <100%> (+0.06%)` | :white_check_mark: |
| [resdk/resources/base.py](https://codecov.io/gh/genialis/resolwe-bio-py/compare/afe16da59e19507196a3591650a71f94394ffe2e...194f027f7b6bacf730c5b5b2fa4bf6795f06bdde?src=pr&el=tree#diff-cmVzZGsvcmVzb3VyY2VzL2Jhc2UucHk=) | `94.89% <20%> (-4.04%)` | :x: |
| [resdk/query.py](https://codecov.io/gh/genialis/resolwe-bio-py/compare/afe16da59e19507196a3591650a71f94394ffe2e...194f027f7b6bacf730c5b5b2fa4bf6795f06bdde?src=pr&el=tree#diff-cmVzZGsvcXVlcnkucHk=) | `88.19% <64.28%> (-4.55%)` | :x: |
| ... and [1 more](https://codecov.io/gh/genialis/resolwe-bio-py/pull/91?src=pr&el=tree-more) | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/genialis/resolwe-bio-py/pull/91?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/genialis/resolwe-bio-py/pull/91?src=pr&el=footer). Last update [afe16da...194f027](https://codecov.io/gh/genialis/resolwe-bio-py/compare/afe16da59e19507196a3591650a71f94394ffe2e...194f027f7b6bacf730c5b5b2fa4bf6795f06bdde?src=pr&el=footer&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).

Patch:
diff --git a/docs/CHANGELOG.rst b/docs/CHANGELOG.rst
index e92a0b2..3aaeaef 100644
--- a/docs/CHANGELOG.rst
+++ b/docs/CHANGELOG.rst
@@ -12,11 +12,22 @@ Added
 -----
 * Option to set API url with ``RESOLWE_HOST_URL`` environment varaible
 
+Added
+-----
+* ``count``, ``delete`` and ``create`` methods to query
+* Support downloading ``basic:dir:`` fields
+
 Changed
 -------
 * Remove ``presample`` endpoint, as it doesn't exist in resolwe anymore
 * Update the way to mark ``sample`` as annotated
+* Add confirmation before deleting an object
 
+Fixed
+-----
+* Fix related queries (i.e. ``collection.data``, ``collection.samples``...)
+  for newly created objects and raise error if they are accessed before object
+  is saved
 
 ==================
 1.6.4 - 2017-02-17
diff --git a/resdk/query.py b/resdk/query.py
index 71b375b..1b38fd0 100644
--- a/resdk/query.py
+++ b/resdk/query.py
@@ -110,6 +110,7 @@ class ResolweQuery(object):
     """
 
     _cache = None
+    _count = None  # number of objects in current query (without applied limit and offset)
     _limit = None
     _offset = None
     _filters = collections.defaultdict(list)
@@ -183,8 +184,7 @@ class ResolweQuery(object):
 
     def __len__(self):
         """Return length of results of current query."""
-        self._fetch()
-        return len(self._cache)
+        return self.count()
 
     def _clone(self):
         """Return copy of current object with empty cache."""
@@ -237,6 +237,7 @@ class ResolweQuery(object):
 
         # Extract data from paginated response
         if isinstance(items, dict) and 'results' in items:
+            self._count = items['count']
             items = items['results']
 
         self._cache = [self._populate_resource(data) for data in items]
@@ -244,6 +245,23 @@ class ResolweQuery(object):
     def clear_cache(self):
         """Clear cache."""
         self._cache = None
+        self._count = None
+
+    def count(self):
+        """Return number of objects in current query."""
+        # pylint: disable=protected-access
+        if self._count is None:
+            count_query = self._clone()
+            count_query._offset = 0
+            count_query._limit = 1
+            count_query._fetch()
+            self._count = count_query._count
+
+        if self._limit is None:
+            return self._count
+
+        remaining = self._count - self._offset
+        return max(0, min(self._limit, remaining))
 
     def get(self, *args, **kwargs):
         """Get object that matches given parameters.
@@ -285,6 +303,10 @@ class ResolweQuery(object):
 
         return response[0]
 
+    def create(self, *args, **kwargs):
+        """Return new instance of current resource."""
+        return self.resource(resolwe=self.resolwe, *args, **kwargs)
+
     def post(self, data):
         """Post data to this endpoint.
 
@@ -298,6 +320,20 @@ class ResolweQuery(object):
         new_query._add_filter(filters)  # pylint: disable=protected-access
         return new_query
 
+    def delete(self, force=False):
+        """Delete objects in current query."""
+        if not force:
+            user_input = six.moves.input(
+                'Do you really want to delete {} object(s)?[yN] '.format(self.count()))
+            if user_input.strip().lower() != 'y':
+                return
+
+        # TODO: Use bulk delete when supported on backend
+        for obj in self:
+            obj.delete(force=True)
+
+        self.clear_cache()
+
     def all(self):
         """Return copy of the current queryset.
 
diff --git a/resdk/resolwe.py b/resdk/resolwe.py
index 187bb38..8b7c110 100644
--- a/resdk/resolwe.py
+++ b/resdk/resolwe.py
@@ -514,11 +514,18 @@ class Resolwe(ResolweUtilsMixin):
 
             for file_uri in files:
                 file_name = os.path.basename(file_uri)
+                file_path = os.path.dirname(file_uri)
                 file_url = urljoin(self.url, 'data/{}'.format(file_uri))
 
-                self.logger.info("* %s", file_name)
+                # Remove data id from path
+                file_path = file_path.split('/', 1)[1] if '/' in file_path else ''
+                full_path = os.path.join(download_dir, file_path)
+                if not os.path.isdir(full_path):
+                    os.makedirs(full_path)
 
-                with open(os.path.join(download_dir, file_name), 'wb') as file_handle:
+                self.logger.info("* %s", os.path.join(file_path, file_name))
+
+                with open(os.path.join(download_dir, file_path, file_name), 'wb') as file_handle:
                     response = requests.get(file_url, stream=True, auth=self.auth)
 
                     if not response.ok:
diff --git a/resdk/resources/base.py b/resdk/resources/base.py
index 2931460..97098db 100644
--- a/resdk/resources/base.py
+++ b/resdk/resources/base.py
@@ -153,8 +153,14 @@ class BaseResource(object):
             response = self.api.post(payload)
             self._update_fields(response)
 
-    def delete(self):
+    def delete(self, force=False):
         """Delete the resource object from the server."""
+        if not force:
+            user_input = six.moves.input('Do you really want to delete {}?[yN] '.format(self))
+
+            if user_input.strip().lower() != 'y':
+                return
+
         self.api(self.id).delete()
 
     def __setattr__(self, name, value):
diff --git a/resdk/resources/collection.py b/resdk/resources/collection.py
index 156c9a5..f75246b 100644
--- a/resdk/resources/collection.py
+++ b/resdk/resources/collection.py
@@ -24,6 +24,9 @@ class BaseCollection(BaseResource):
 
     """
 
+    #: lazy loaded list of data objects
+    _data = None
+
     WRITABLE_FIELDS = ('description', 'settings', 'descriptor_schema',
                        'descriptor') + BaseResource.WRITABLE_FIELDS
 
@@ -38,14 +41,17 @@ class BaseCollection(BaseResource):
         self.descriptor = None
         #: descriptor schema
         self.descriptor_schema = None
-        #: list of data objects - SET THIS IN SUBCLASS
-        self.data = []
 
         super(BaseCollection, self).__init__(slug, id, model_data, resolwe)
 
+    @property
+    def data(self):
+        """Return list of attached Data objects."""
+        raise NotImplementedError('This should be implemented in subclass')
+
     def _clear_data_cache(self):
         """Clear data cache."""
-        raise NotImplementedError('This should be implemented in subclass')
+        self._data = None
 
     def add_data(self, *data):
         """Add ``data`` objects to the collection."""
@@ -134,25 +140,40 @@ class Collection(BaseCollection):
 
     endpoint = 'collection'
 
+    #: (lazy loaded) list of samples that belong to collection
+    _samples = None
+
     def __init__(self, slug=None, id=None,  # pylint: disable=redefined-builtin
                  model_data=None, resolwe=None):
         """Initialize attributes."""
         BaseCollection.__init__(self, slug, id, model_data, resolwe)
 
-        #: (lazy loaded) list of data object that belong to collection
-        self.data = self.resolwe.data.filter(collection=self.id)
-        #: (lazy loaded) list of samples that belong to collection
-        self.samples = self.resolwe.sample.filter(collections=self.id)
-
     def update(self):
         """Clear cache and update resource fields from the server."""
-        self.data.clear_cache()  # pylint: disable=no-member
-        self.samples.clear_cache()
+        self._data = None
+        self._samples = None
 
         super(Collection, self).update()
 
-    def _clear_data_cache(self):
-        self.data.clear_cache()  # pylint: disable=no-member
+    @property
+    def data(self):
+        """Return list of data objects on collection."""
+        if self.id is None:
+            raise ValueError('Instance must be saved before accessing `data` attribute.')
+        if self._data is None:
+            self._data = self.resolwe.data.filter(collection=self.id)
+
+        return self._data
+
+    @property
+    def samples(self):
+        """Return list of samples on collection."""
+        if self.id is None:
+            raise ValueError('Instance must be saved before accessing `samples` attribute.')
+        if self._samples is None:
+            self._samples = self.resolwe.sample.filter(collections=self.id)
+
+        return self._samples
 
     def add_samples(self, *samples):
         """Add `samples` objects to the collection."""
diff --git a/resdk/resources/data.py b/resdk/resources/data.py
index 1e3fff2..dedd68f 100644
--- a/resdk/resources/data.py
+++ b/resdk/resources/data.py
@@ -1,6 +1,7 @@
 """Data resource."""
 from __future__ import absolute_import, division, print_function, unicode_literals
 
+import json
 import logging
 
 import requests
@@ -32,6 +33,8 @@ class Data(BaseResource):
 
     #: (lazy loaded) annotated ``Sample`` to which ``Data`` object belongs
     _sample = None
+    #: (lazy loaded) list of collections to which data object belongs
+    _collections = None
 
     WRITABLE_FIELDS = ('descriptor_schema', 'descriptor') + BaseResource.WRITABLE_FIELDS
     UPDATE_PROTECTED_FIELDS = ('input', 'process') + BaseResource.UPDATE_PROTECTED_FIELDS
@@ -83,15 +86,12 @@ class Data(BaseResource):
 
         super(Data, self).__init__(slug, id, model_data, resolwe)
 
-        #: (lazy loaded) list of collections to which data object belongs
-        self.collections = self.resolwe.collection.filter(data=self.id)
-
         self.logger = logging.getLogger(__name__)
 
     def update(self):
         """Clear cache and update resource fields from the server."""
         self._sample = None
-        self.collections.clear_cache()
+        self._collections = None
 
         super(Data, self).update()
 
@@ -137,35 +137,37 @@ class Data(BaseResource):
 
         return flat
 
+    @property
+    def collections(self):
+        """Return list of collections to which data object belongs."""
+        if self.id is None:
+            raise ValueError('Instance must be saved before accessing `collections` attribute.')
+        if self._collections is None:
+            self._collections = self.resolwe.collection.filter(data=self.id)
+
+        return self._collections
+
     @property
     def sample(self):
         """Get ``sample`` that object belongs to."""
+        if self.id is None:
+            raise ValueError('Instance must be saved before accessing `sample` attribute.')
         if self._sample is None:
             self._sample = self.resolwe.sample.filter(data=self.id)
             self._sample = None if len(self._sample) == 0 else self._sample[0]
         return self._sample
 
-    def files(self, file_name=None, field_name=None):
-        """Get list of downloadable fields.
-
-        Filter files by file name or output field.
-
-        :param file_name: name of file
-        :type file_name: string
-        :param field_name: output field name
-        :type field_name: string
-        :rtype: List of tuples (data_id, file_name, field_name, process_type)
-
-        """
+    def _files_dirs(self, field_type, file_name=None, field_name=None):
+        """Get list of downloadable fields."""
         download_list = []
 
         def put_in_download_list(elm, fname):
-            """Append only files with equal name."""
-            if 'file' in elm:
-                if file_name is None or file_name == elm['file']:
-                    download_list.append(elm['file'])
+            """Append only files od dirs with equal name."""
+            if field_type in elm:
+                if file_name is None or file_name == elm[field_type]:
+                    download_list.append(elm[field_type])
             else:
-                raise KeyError("Item {} does not contain 'file' key.".format(fname))
+                raise KeyError("Item {} does not contain '{}' key.".format(fname, field_type))
 
         if field_name and not field_name.startswith('output.'):
             field_name = 'output.{}'.format(field_name)
@@ -174,30 +176,76 @@ class Data(BaseResource):
             if (ann_field_name.startswith('output') and
                     (field_name is None or field_name == ann_field_name) and
                     ann['value'] is not None):
-                if ann['type'].startswith('basic:file:'):
+                if ann['type'].startswith('basic:{}:'.format(field_type)):
                     put_in_download_list(ann['value'], ann_field_name)
-                elif ann['type'].startswith('list:basic:file:'):
+                elif ann['type'].startswith('list:basic:{}:'.format(field_type)):
                     for element in ann['value']:
                         put_in_download_list(element, ann_field_name)
 
         return download_list
 
-    def download(self, file_name=None, field_name=None, download_dir=None):
-        """Download Data object's files.
+    def _get_dir_files(self, dir_name):
+        files_list, dir_list = [], []
+
+        dir_url = urljoin(self.resolwe.url, 'data/{}/{}'.format(self.id, dir_name))
+        if not dir_url.endswith('/'):
+            dir_url += '/'
+        response = requests.get(dir_url, auth=self.resolwe.auth)
+        response = json.loads(response.content.decode('utf-8'))
+
+        for obj in response:
+            obj_path = '{}/{}'.format(dir_name, obj['name'])
+            if obj['type'] == 'directory':
+                dir_list.append(obj_path)
+            else:
+                files_list.append(obj_path)
 
-        Download files from the Resolwe server to the download
-        directory (defaults to the current working directory).
+        if dir_list:
+            for new_dir in dir_list:
+                files_list.extend(self._get_dir_files(new_dir))
+
+        return files_list
+
+    def files(self, file_name=None, field_name=None):
+        """Get list of downloadable file fields.
+
+        Filter files by file name or output field.
 
         :param file_name: name of file
         :type file_name: string
-        :param field_name: file field name
+        :param field_name: output field name
+        :type field_name: string
+        :rtype: List of tuples (data_id, file_name, field_name, process_type)
+
+        """
+        if not self.id:
+            raise ValueError('Instance must be saved before using `files` method.')
+
+        file_list = self._files_dirs('file', file_name, field_name)
+
+        for dir_name in self._files_dirs('dir', file_name, field_name):
+            print('>>', self._get_dir_files)
+            file_list.extend(self._get_dir_files(dir_name))
+
+        return file_list
+
+    def download(self, file_name=None, field_name=None, download_dir=None):
+        """Download Data object's files and directories.
+
+        Download files and directoriesfrom the Resolwe server to the
+        download directory (defaults to the current working directory).
+
+        :param file_name: name of file or directory
+        :type file_name: string
+        :param field_name: file or directory field name
         :type field_name: string
         :param download_dir: download path
         :type download_dir: string
         :rtype: None
 
-        Data objects can contain multiple files. All are downloaded by
-        default, but may be filtered by name or output field:
+        Data objects can contain multiple files and directories. All are
+        downloaded by default, but may be filtered by name or output
+        field:
 
         * re.data.get(42).download(file_name='alignment7.bam')
         * re.data.get(42).download(field_name='bam')
diff --git a/resdk/resources/sample.py b/resdk/resources/sample.py
index 016205d..0c785b9 100644
--- a/resdk/resources/sample.py
+++ b/resdk/resources/sample.py
@@ -25,25 +25,40 @@ class Sample(SampleUtilsMixin, BaseCollection):
 
     endpoint = 'sample'
 
+    #: (lazy loaded) list of collections  to which object belongs
+    _collections = None
+
     def __init__(self, slug=None, id=None, model_data=None,  # pylint: disable=redefined-builtin
                  resolwe=None):
         """Initialize attributes."""
         super(Sample, self).__init__(slug, id, model_data, resolwe)
 
-        #: (lazy loaded) list of collections  to which object belongs
-        self.collections = self.resolwe.collection.filter(entity=self.id)
-        #: (lazy loaded) list of data object that belong to sample
-        self.data = self.resolwe.data.filter(entity=self.id)
-
     def update(self):
         """Clear cache and update resource fields from the server."""
-        self.collections.clear_cache()
-        self.data.clear_cache()  # pylint: disable=no-member
+        self._collections = None
+        self._data = None
 
         super(Sample, self).update()
 
-    def _clear_data_cache(self):
-        self.data.clear_cache()  # pylint: disable=no-member
+    @property
+    def data(self):
+        """Return list of data objects on collection."""
+        if self.id is None:
+            raise ValueError('Instance must be saved before accessing `data` attribute.')
+        if self._data is None:
+            self._data = self.resolwe.data.filter(entity=self.id)
+
+        return self._data
+
+    @property
+    def collections(self):
+        """Return list of collections to which sample belongs."""
+        if self.id is None:
+            raise ValueError('Instance must be saved before accessing `collections` attribute.')
+        if self._collections is None:
+            self._collections = self.resolwe.collection.filter(entity=self.id)
+
+        return self._collections
 
     def print_annotation(self):
         """Provide annotation data."""


Test patch:
diff --git a/resdk/tests/unit/test_collections.py b/resdk/tests/unit/test_collections.py
index 57d8ddd..e35dcbb 100644
--- a/resdk/tests/unit/test_collections.py
+++ b/resdk/tests/unit/test_collections.py
@@ -75,52 +75,88 @@ class TestCollection(unittest.TestCase):
             Collection.print_annotation(collection_mock)
 
     def test_data(self):
-        resolwe_mock = MagicMock(**{'data.filter.return_value': ['data_1', 'data_2', 'data_3']})
-        collection = Collection(id=1, resolwe=resolwe_mock)
+        collection = Collection(id=1, resolwe=MagicMock())
 
+        # test getting data attribute
+        collection.resolwe.data.filter = MagicMock(return_value=['data_1', 'data_2', 'data_3'])
         self.assertEqual(collection.data, ['data_1', 'data_2', 'data_3'])
 
+        # test caching data attribute
+        self.assertEqual(collection.data, ['data_1', 'data_2', 'data_3'])
+        self.assertEqual(collection.resolwe.data.filter.call_count, 1)
+
         # cache is cleared at update
-        collection.data = MagicMock()
+        collection._data = ['data']
         collection.update()
-        self.assertEqual(collection.data.clear_cache.call_count, 1)
+        self.assertEqual(collection._data, None)
+
+        # raising error if collection is not saved
+        collection.id = None
+        with self.assertRaises(ValueError):
+            _ = collection.data
 
     def test_samples(self):
-        resolwe_mock = MagicMock(**{'sample.filter.return_value': ['sample1', 'sample2']})
-        collection = Collection(id=1, resolwe=resolwe_mock)
+        collection = Collection(id=1, resolwe=MagicMock())
 
+        # test getting samples attribute
+        collection.resolwe.sample.filter = MagicMock(return_value=['sample1', 'sample2'])
         self.assertEqual(collection.samples, ['sample1', 'sample2'])
 
         # cache is cleared at update
-        collection.samples = MagicMock()
+        collection._samples = ['sample']
         collection.update()
-        self.assertEqual(collection.samples.clear_cache.call_count, 1)
+        self.assertEqual(collection._samples, None)
+
+        # raising error if data collection is not saved
+        collection.id = None
+        with self.assertRaises(ValueError):
+            _ = collection.samples
 
 
 class TestSample(unittest.TestCase):
 
     def test_data(self):
-        resolwe_mock = MagicMock(**{'data.filter.return_value': ['data_1', 'data_2', 'data_3']})
-        sample = Sample(id=1, resolwe=resolwe_mock)
+        sample = Sample(id=1, resolwe=MagicMock())
 
+        # test getting data attribute
+        sample.resolwe.data.filter = MagicMock(return_value=['data_1', 'data_2', 'data_3'])
         self.assertEqual(sample.data, ['data_1', 'data_2', 'data_3'])
 
+        # test caching data attribute
+        self.assertEqual(sample.data, ['data_1', 'data_2', 'data_3'])
+        self.assertEqual(sample.resolwe.data.filter.call_count, 1)
+
         # cache is cleared at update
-        sample.data = MagicMock()
+        sample._data = ['data']
         sample.update()
-        self.assertEqual(sample.data.clear_cache.call_count, 1)  # pylint: disable=no-member
+        self.assertEqual(sample._data, None)
+
+        # raising error if sample is not saved
+        sample.id = None
+        with self.assertRaises(ValueError):
+            _ = sample.data
 
     def test_collections(self):
-        resolwe_mock = MagicMock(
-            **{'collection.filter.return_value': ['collection_1', 'collection_2']})
-        sample = Sample(id=1, resolwe=resolwe_mock)
+        sample = Sample(id=1, resolwe=MagicMock())
 
-        self.assertEqual(sample.collections, ['collection_1', 'collection_2'])
+        # test getting data attribute
+        sample.resolwe.collection.filter = MagicMock(
+            return_value=['collection_1', 'collection_2', 'collection_3'])
+        self.assertEqual(sample.collections, ['collection_1', 'collection_2', 'collection_3'])
+
+        # test caching data attribute
+        self.assertEqual(sample.collections, ['collection_1', 'collection_2', 'collection_3'])
+        self.assertEqual(sample.resolwe.collection.filter.call_count, 1)
 
         # cache is cleared at update
-        sample.collections = MagicMock()
+        sample._collections = ['collection']
         sample.update()
-        self.assertEqual(sample.collections.clear_cache.call_count, 1)
+        self.assertEqual(sample._collections, None)
+
+        # raising error if sample is not saved
+        sample.id = None
+        with self.assertRaises(ValueError):
+            _ = sample.collections
 
     @patch('resdk.resources.sample.Sample', spec=True)
     def test_sample_print_annotation(self, sample_mock):
diff --git a/resdk/tests/unit/test_data.py b/resdk/tests/unit/test_data.py
index c590ab6..d0fd76c 100644
--- a/resdk/tests/unit/test_data.py
+++ b/resdk/tests/unit/test_data.py
@@ -2,6 +2,7 @@
 Unit tests for resdk/resources/data.py file.
 """
 # pylint: disable=missing-docstring, protected-access
+from __future__ import absolute_import, division, print_function, unicode_literals
 
 import unittest
 
@@ -49,44 +50,83 @@ class TestData(unittest.TestCase):
         data.update()
         self.assertEqual(data._sample, None)
 
+        # raising error if data object is not saved
+        data.id = None
+        with self.assertRaises(ValueError):
+            _ = data.sample
+
     def test_collections(self):
-        resolwe_mock = MagicMock(
-            **{'collection.filter.return_value': ['collection_1', 'collection_2']})
-        data = Data(id=1, resolwe=resolwe_mock)
+        data = Data(id=1, resolwe=MagicMock())
 
-        self.assertEqual(data.collections, ['collection_1', 'collection_2'])
+        # test getting collections attribute
+        data.resolwe.collection.filter = MagicMock(return_value=['collection'])
+        self.assertEqual(data.collections, ['collection'])
+
+        # test caching collections attribute
+        self.assertEqual(data.collections, ['collection'])
+        self.assertEqual(data.resolwe.collection.filter.call_count, 1)
 
         # cache is cleared at update
-        data.collections = MagicMock()
+        data._collections = ['collection']
         data.update()
-        self.assertEqual(data.collections.clear_cache.call_count, 1)
+        self.assertEqual(data._collections, None)
 
-    @patch('resdk.resources.data.Data', spec=True)
-    def test_files(self, data_mock):
-        data_annotation = {
+        # raising error if data object is not saved
+        data.id = None
+        with self.assertRaises(ValueError):
+            _ = data.collections
+
+    def test_files(self):
+        data = Data(id=123, resolwe=MagicMock())
+        data._get_dir_files = MagicMock(
+            side_effect=[['first_dir/file1.txt'], ['fastq_dir/file2.txt']])
+
+        data.annotation = {
             'output.list': {'value': [{'file': "element.gz"}], 'type': 'list:basic:file:'},
+            'output.dir_list': {'value': [{'dir': "first_dir"}], 'type': 'list:basic:dir:'},
             'output.fastq': {'value': {'file': "file.fastq.gz"}, 'type': 'basic:file:fastq'},
             'output.fastq_archive': {'value': {'file': "archive.gz"}, 'type': 'basic:file:'},
+            'output.fastq_dir': {'value': {'dir': "fastq_dir"}, 'type': 'basic:dir:'},
             'input.fastq_url': {'value': {'file': "blah"}, 'type': 'basic:url:'},
             'input.blah': {'value': "blah.gz", 'type': 'basic:file:'}
         }
-        bad_data_annotation = {
+
+        file_list = data.files()
+        six.assertCountEqual(self, file_list, [
+            'element.gz',
+            'archive.gz',
+            'file.fastq.gz',
+            'first_dir/file1.txt',
+            'fastq_dir/file2.txt'
+        ])
+        file_list = data.files(file_name='element.gz')
+        self.assertEqual(file_list, ['element.gz'])
+        file_list = data.files(field_name='output.fastq')
+        self.assertEqual(file_list, ['file.fastq.gz'])
+
+        data.annotation = {
             'output.list': {'value': [{'no_file_field_here': "element.gz"}],
                             'type': 'list:basic:file:'},
         }
-        data_mock.configure_mock(annotation=data_annotation)
+        with six.assertRaisesRegex(self, KeyError, "does not contain 'file' key."):
+            data.files()
 
-        file_list = Data.files(data_mock)
-        self.assertEqual(set(file_list), set(['element.gz', 'archive.gz', 'file.fastq.gz']))
-        file_list = Data.files(data_mock, file_name='element.gz')
-        self.assertEqual(file_list, ['element.gz'])
-        file_list = Data.files(data_mock, field_name='output.fastq')
-        self.assertEqual(file_list, ['file.fastq.gz'])
+        data = Data(resolwe=MagicMock())
+        with six.assertRaisesRegex(self, ValueError, "must be saved before"):
+            data.files()
+
+    @patch('resdk.resources.data.requests')
+    def test_dir_files(self, requests_mock):
+        data = Data(id=123, resolwe=MagicMock(url='http://resolwe.url'))
+        requests_mock.get = MagicMock(side_effect=[
+            MagicMock(content=b'[{"type": "file", "name": "file1.txt"}, '
+                              b'{"type": "directory", "name": "subdir"}]'),
+            MagicMock(content=b'[{"type": "file", "name": "file2.txt"}]'),
+        ])
+
+        files = data._get_dir_files('test_dir')
 
-        data_mock.configure_mock(annotation=bad_data_annotation)
-        message = r"Item .* does not contain 'file' key."
-        with six.assertRaisesRegex(self, KeyError, message):
-            Data.files(data_mock)
+        self.assertEqual(files, ['test_dir/file1.txt', 'test_dir/subdir/file2.txt'])
 
     @patch('resdk.resources.data.Data', spec=True)
     def test_download_fail(self, data_mock):
@@ -97,16 +137,18 @@ class TestData(unittest.TestCase):
     @patch('resdk.resources.data.Data', spec=True)
     def test_download_ok(self, data_mock):
         data_mock.configure_mock(id=123, **{'resolwe': MagicMock()})
-        data_mock.configure_mock(**{'files.return_value': ['file1.txt', 'file2.fq.gz']})
+        data_mock.configure_mock(**{
+            'files.return_value': ['file1.txt', 'file2.fq.gz'],
+        })
 
         Data.download(data_mock)
         data_mock.resolwe._download_files.assert_called_once_with(
-            [u'123/file1.txt', u'123/file2.fq.gz'], None)
+            ['123/file1.txt', '123/file2.fq.gz'], None)
 
         data_mock.reset_mock()
         Data.download(data_mock, download_dir="/some/path/")
         data_mock.resolwe._download_files.assert_called_once_with(
-            [u'123/file1.txt', u'123/file2.fq.gz'], '/some/path/')
+            ['123/file1.txt', '123/file2.fq.gz'], '/some/path/')
 
     @patch('resdk.resources.data.Data', spec=True)
     def test_add_output(self, data_mock):
@@ -115,10 +157,10 @@ class TestData(unittest.TestCase):
                         'output.fasta': {'type': 'basic:file:', 'value': {'file': 'genome.fa'}}}
         )
 
-        files_list = Data.files(data_mock, field_name="output.fastq")
+        files_list = Data._files_dirs(data_mock, 'file', field_name="output.fastq")
         self.assertEqual(files_list, ['reads.fq'])
 
-        files_list = Data.files(data_mock, field_name="fastq")
+        files_list = Data._files_dirs(data_mock, 'file', field_name="fastq")
         self.assertEqual(files_list, ['reads.fq'])
 
     @patch('resdk.resources.data.Data', spec=True)
diff --git a/resdk/tests/unit/test_query.py b/resdk/tests/unit/test_query.py
index 74523b7..df38837 100644
--- a/resdk/tests/unit/test_query.py
+++ b/resdk/tests/unit/test_query.py
@@ -103,7 +103,7 @@ class TestResolweQuery(unittest.TestCase):
         self.assertEqual(rep, '[1,\n 2,\n 3]')
 
     def test_len(self):
-        query = MagicMock(spec=ResolweQuery, _cache=[1, 2, 3])
+        query = MagicMock(spec=ResolweQuery, **{'count.return_value': 3})
         query.__len__ = ResolweQuery.__len__
 
         self.assertEqual(len(query), 3)
@@ -166,6 +166,30 @@ class TestResolweQuery(unittest.TestCase):
         ResolweQuery.clear_cache(query)
         self.assertEqual(query._cache, None)
 
+    def test_count(self):
+        count_query = MagicMock(spec=ResolweQuery, _count=10)
+        query = MagicMock(spec=ResolweQuery, _count=None, _limit=None, _offset=None,
+                          **{'_clone.return_value': count_query})
+
+        self.assertEqual(ResolweQuery.count(query), 10)
+
+        query._limit = 2
+        query._offset = 0
+        self.assertEqual(ResolweQuery.count(query), 2)
+
+        query._limit = 2
+        query._offset = 9
+        self.assertEqual(ResolweQuery.count(query), 1)
+
+        query._limit = 2
+        query._offset = 12
+        self.assertEqual(ResolweQuery.count(query), 0)
+
+        query._count = 5
+        query._limit = None
+        query._offset = None
+        self.assertEqual(ResolweQuery.count(query), 5)
+
     def test_get(self):
         new_query = MagicMock(spec=ResolweQuery)
         query = MagicMock(spec=ResolweQuery, **{'_clone.return_value': new_query})
diff --git a/resdk/tests/unit/test_resolwe.py b/resdk/tests/unit/test_resolwe.py
index fb99b6d..a722215 100644
--- a/resdk/tests/unit/test_resolwe.py
+++ b/resdk/tests/unit/test_resolwe.py
@@ -569,7 +569,6 @@ class TestDownload(unittest.TestCase):
     @patch('resdk.resolwe.Resolwe', spec=True)
     def test_empty_file_list(self, resolwe_mock, os_mock):
         resolwe_mock.configure_mock(**self.config)
-        os_mock.path.isfile.return_value = True
 
         Resolwe._download_files(resolwe_mock, [])