Address comments

E3SM-Project · Dec 6, 2021 · e139de7 · e139de7
1 parent baffaaa
commit e139de7
Show file tree

Hide file tree

Showing 3 changed files with 51 additions and 3 deletions.
diff --git a/tests/test_check.py b/tests/test_check.py
@@ -232,6 +232,7 @@ def helperCheckTars(
 
     tester.assertWorkspace()
     os.chdir(tester.test_dir)
+
     # Starting at 000001 until the end
     zstash_cmd = '{}zstash check --hpss={} --tars="000001-"{}'.format(
         zstash_path, tester.hpss_path, worker_str
@@ -333,6 +334,22 @@ def helperCheckTars(
         "INFO: Opening tar archive zstash/000007.tar",
     ]
     tester.check_strings(zstash_cmd, output + err, expected_present, expected_absent)
+    # Ending with nonexistent tar
+    zstash_cmd = '{}zstash check --hpss={} --tars="000001-00003"{} file'.format(
+        zstash_path, tester.hpss_path, worker_str
+    )
+    output, err = run_cmd(zstash_cmd)
+    expected_present = ["ValueError: If --tars is used, <files> should not be listed."]
+    expected_absent = [
+        "INFO: Opening tar archive zstash/000000.tar",
+        "INFO: Opening tar archive zstash/000001.tar",
+        "INFO: Opening tar archive zstash/000002.tar",
+        "INFO: Opening tar archive zstash/000003.tar",
+        "INFO: Opening tar archive zstash/000004.tar",
+        "INFO: No failures detected when checking the files.",
+    ]
+    tester.check_strings(zstash_cmd, output + err, expected_present, expected_absent)
+
     os.chdir(TOP_LEVEL)
 
 

diff --git a/tests/test_functions.py b/tests/test_functions.py
@@ -49,6 +49,21 @@ def testParseTarsOption(self):
             ],
         )
 
+        # Check removal of duplicates and sorting
+        tar_list = parse_tars_option("000009,000003,-000005", "000000", "00005d")
+        self.assertEqual(
+            tar_list,
+            ["000000", "000001", "000002", "000003", "000004", "000005", "000009"],
+        )
+
+        # Remove .tar suffix
+        tar_list = parse_tars_option(
+            "000009.tar-00000a,000003.tar,-000002.tar", "000000", "00005d"
+        )
+        self.assertEqual(
+            tar_list, ["000000", "000001", "000002", "000003", "000009", "00000a"]
+        )
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/zstash/extract.py b/zstash/extract.py
@@ -123,12 +123,25 @@ def parse_tars_option(tars: str, first_tar: str, last_tar: str) -> List[str]:
         m: Optional[re.Match]
         m = re.match("(.*)-(.*)", tar_str)
         if m:
-            beginning_tar = int(m.group(1), 16)
-            ending_tar = int(m.group(2), 16)
+            m1: str = m.group(1)
+            m2: str = m.group(2)
+            # Remove .tar suffix
+            if m1.endswith(".tar"):
+                m1 = m1[:-4]
+            if m2.endswith(".tar"):
+                m2 = m2[:-4]
+            beginning_tar: int = int(m1, 16)
+            ending_tar: int = int(m2, 16)
+            t: int
             for t in range(beginning_tar, ending_tar + 1):
                 tar_list.append("{:06x}".format(t))
         else:
+            # Remove .tar suffix
+            if tar_str.endswith(".tar"):
+                tar_str = tar_str[:-4]
             tar_list.append(tar_str)
+    # Remove duplicates and sort tar_list
+    tar_list = sorted(list(set(tar_list)))
     return tar_list
 
 
@@ -191,6 +204,9 @@ def extract_database(
 
     matches_: List[TupleFilesRow] = []
     if args.tars is not None:
+        # Ignore default value for args.files ("*")
+        if args.files != ["*"]:
+            raise ValueError("If --tars is used, <files> should not be listed.")
         tar_names_initial: List[Tuple[str]] = cur.execute(
             u"select distinct tar from files"
         ).fetchall()
@@ -205,7 +221,7 @@ def extract_database(
                 (tar + ".tar",),
             )
             matches_ = matches_ + cur.fetchall()
-    else:  # TODO: should this always be run (i.e., not if/else?)
+    else:
         # Find matching files
         for args_file in args.files:
             cur.execute(