From 42f4af45aee750307d435b8384ae57fd5d9a9eb6 Mon Sep 17 00:00:00 2001
From: Alistair Coles <alistairncoles@gmail.com>
Date: Mon, 15 Mar 2021 18:16:32 +0000
Subject: [PATCH] Add test coverage for audit_location_generator yield_hashes
 arg

Add unit test coverage for the optional yield_hash_dirs arg
added to audit_location_generator in [1].

[1] Related-Change: I6e01ed29eb8971eeff96887efbfdfb76a01b00c3
Change-Id: Ibd738365112e0efcb9a7dad73cf4a09f1c3236bc
---
 test/unit/common/test_utils.py | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/test/unit/common/test_utils.py b/test/unit/common/test_utils.py
index cf878f6e3b..a9f99ebca0 100644
--- a/test/unit/common/test_utils.py
+++ b/test/unit/common/test_utils.py
@@ -6259,6 +6259,7 @@ class TestAuditLocationGenerator(unittest.TestCase):
     def test_find_objects(self):
         with temptree([]) as tmpdir:
             expected_objs = list()
+            expected_dirs = list()
             logger = FakeLogger()
             data = os.path.join(tmpdir, "drive", "data")
             os.makedirs(data)
@@ -6270,6 +6271,7 @@ class TestAuditLocationGenerator(unittest.TestCase):
             os.makedirs(suffix)
             hash_path = os.path.join(suffix, "hash")
             os.makedirs(hash_path)
+            expected_dirs.append((hash_path, 'drive', 'partition1'))
             obj_path = os.path.join(hash_path, "obj1.db")
             with open(obj_path, "w"):
                 pass
@@ -6280,6 +6282,7 @@ class TestAuditLocationGenerator(unittest.TestCase):
             os.makedirs(suffix)
             hash_path = os.path.join(suffix, "hash2")
             os.makedirs(hash_path)
+            expected_dirs.append((hash_path, 'drive', 'partition2'))
             obj_path = os.path.join(hash_path, "obj2.db")
             with open(obj_path, "w"):
                 pass
@@ -6292,6 +6295,14 @@ class TestAuditLocationGenerator(unittest.TestCase):
             self.assertEqual(sorted(got_objs), sorted(expected_objs))
             self.assertEqual(1, len(logger.get_lines_for_level('warning')))
 
+            # check yield_hash_dirs option
+            locations = utils.audit_location_generator(
+                tmpdir, "data", mount_check=False, logger=logger,
+                yield_hash_dirs=True,
+            )
+            got_dirs = list(locations)
+            self.assertEqual(sorted(got_dirs), sorted(expected_dirs))
+
     def test_ignore_metadata(self):
         with temptree([]) as tmpdir:
             logger = FakeLogger()