From dfd69889803f97a7afb2df7dbc24085517eb12fe Mon Sep 17 00:00:00 2001
From: Arman <30441115+armanbm177@users.noreply.github.com>
Date: Sat, 13 Jan 2024 19:25:11 -0800
Subject: [PATCH] implemented comma separated part numbers for run-tests.py
 (#66)

Co-authored-by: ibm5 <ibm5@cs.washington.edu>
---
 .../testing/junit/DSLabsTestCore.java         | 34 +++++++++++++------
 handout-files/run-tests.py                    |  5 +--
 2 files changed, 26 insertions(+), 13 deletions(-)

diff --git a/framework/tst/dslabs/framework/testing/junit/DSLabsTestCore.java b/framework/tst/dslabs/framework/testing/junit/DSLabsTestCore.java
index a603bdf0..04e1a36f 100644
--- a/framework/tst/dslabs/framework/testing/junit/DSLabsTestCore.java
+++ b/framework/tst/dslabs/framework/testing/junit/DSLabsTestCore.java
@@ -23,6 +23,7 @@
 package dslabs.framework.testing.junit;
 
 import com.google.common.base.Throwables;
+import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Sets;
 import dslabs.framework.testing.utils.ClassSearch;
 import java.util.Set;
@@ -121,10 +122,10 @@ public static void main(String[] args) throws Exception {
         Option.builder("p")
             .longOpt("part")
             .required(false)
-            .argName("PART")
             .hasArg(true)
             .numberOfArgs(1)
-            .desc("part number")
+            .argName("PART_NUMS")
+            .desc("comma-separated list of part numbers")
             .build();
     final Option testNum =
         Option.builder("n")
@@ -196,9 +197,13 @@ public String describe() {
               }
             });
 
+    final ImmutableSet<String> partNumbers =
+        line.hasOption(part)
+            ? ImmutableSet.copyOf(line.getOptionValue(part).split(","))
+            : ImmutableSet.of();
+
     // Only run test classes for this part
-    if (line.hasOption(part)) {
-      final int partNum = Integer.parseInt(line.getOptionValue(part));
+    if (!partNumbers.isEmpty()) {
       request =
           request.filterWith(
               new Filter() {
@@ -208,25 +213,32 @@ public boolean shouldRun(Description description) {
                     return true;
                   }
                   var p = description.getAnnotation(Part.class);
-                  return p != null && p.value() == partNum;
+                  return p != null && partNumbers.contains(String.valueOf(p.value()));
                 }
 
                 @Override
                 public String describe() {
-                  return "part " + partNum;
+                  return String.format("part %s", partNumbers);
                 }
               });
     }
 
     // Only run chosen test numbers
     if (line.hasOption(testNum)) {
-      final Set<String> testNumbers = Sets.newHashSet(line.getOptionValue(testNum).split(","));
+      final ImmutableSet<String> specifiedTestNumbers =
+          ImmutableSet.copyOf(line.getOptionValue(testNum).split(","));
+      final Set<String> testNumbers = Sets.newHashSet(specifiedTestNumbers);
 
       // If the part is selected, allow specifying test number only
-      if (line.hasOption(part)) {
-        for (var tn : line.getOptionValue(testNum).split(",")) {
+      if (!partNumbers.isEmpty()) {
+        if (partNumbers.size() > 1) {
+          throw new IllegalArgumentException(
+              "Cannot specify multiple part numbers when selecting individual test numbers.");
+        }
+        final String partStr = partNumbers.stream().findFirst().get();
+        for (var tn : specifiedTestNumbers) {
           // XXX: not the best way of doing this, but it works
-          testNumbers.add(line.getOptionValue(part) + "." + tn);
+          testNumbers.add(partStr + "." + tn);
         }
       }
 
@@ -243,7 +255,7 @@ public boolean shouldRun(Description description) {
 
                 @Override
                 public String describe() {
-                  return String.format("test numbers %s", testNumbers);
+                  return String.format("test numbers %s", specifiedTestNumbers);
                 }
               });
     }
diff --git a/handout-files/run-tests.py b/handout-files/run-tests.py
index 5372c0c9..0fcba024 100755
--- a/handout-files/run-tests.py
+++ b/handout-files/run-tests.py
@@ -199,8 +199,9 @@ def main():
     group = run_modes.add_mutually_exclusive_group()
 
     parser.add_argument('-l', '--lab', help="lab to run tests for")
-    parser.add_argument('-p', '--part', type=int,
-                        help="part number for tests to run")
+
+    parser.add_argument('-p', '--part',
+                        help="part number(s) to run tests for (comma-separated)")
 
     parser.add_argument('-n', '--test-num',
                         help="specific, comma-separated test numbers to run "