Repository: troycomi/reportseff

Issue:
`MemEff` = 0% if I add argument `-s TIMEOUT`
Version 2.7.6

> reportseff -u $U

```
    JobID    State       Elapsed  TimeEff   CPUEff   MemEff 
  xxxxxxx   TIMEOUT     17:00:06  100.0%     5.0%    72.2%  
  xxxxxxx   TIMEOUT     15:00:18  100.0%    16.0%    74.5%  
...
```

Good. But if I run:

> reportseff -u $U -s TIMEOUT

Then all the `MemEff` values are zeroed, at least for `TIMEOUT` jobs:

```
    JobID   State      Elapsed  TimeEff   CPUEff   MemEff 
  xxxxxxx  TIMEOUT    17:00:06  100.0%     5.0%     0.0%  
  xxxxxxx  TIMEOUT    15:00:18  100.0%    16.0%     0.0%  
...
  ```

Discussion:

Patch:
diff --git a/pyproject.toml b/pyproject.toml
index 7eb03ca..3ddd09a 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "reportseff"
-version = "2.7.6"
+version = "2.7.7"
 description= "Tablular seff output"
 authors = ["Troy Comi <troycomi@gmail.com>"]
 license = "MIT"
diff --git a/src/reportseff/db_inquirer.py b/src/reportseff/db_inquirer.py
index 3e91e46..f40d9f4 100644
--- a/src/reportseff/db_inquirer.py
+++ b/src/reportseff/db_inquirer.py
@@ -266,13 +266,26 @@ class SacctInquirer(BaseInquirer):
         sacct_split = re.compile(r"\^\|\^")
         result = [dict(zip(columns, sacct_split.split(line))) for line in lines if line]
 
-        if self.state:
-            # split to get first word in entries like "CANCELLED BY X"
-            result = [r for r in result if r["State"].split()[0] in self.state]
-
-        if self.not_state:
-            # split to get first word in entries like "CANCELLED BY X"
-            result = [r for r in result if r["State"].split()[0] not in self.not_state]
+        # Sometimes the main job has a different state than the sub jobs
+        # e.g. timeouts have a state of canceled for the batch jobs.
+        # When state filtering is active, need to filter main ids, then retain
+        # only the jobs with matching job ids
+        if self.state or self.not_state:
+            main_jobs = [r for r in result if "." not in r["JobID"]]
+            if self.state:
+                # split to get first word in entries like "CANCELLED BY X"
+                main_jobs = [
+                    r for r in main_jobs if r["State"].split()[0] in self.state
+                ]
+
+            if self.not_state:
+                # split to get first word in entries like "CANCELLED BY X"
+                main_jobs = [
+                    r for r in main_jobs if r["State"].split()[0] not in self.not_state
+                ]
+
+            main_job_ids = {r["JobID"] for r in main_jobs}
+            result = [r for r in result if r["JobID"].split(".")[0] in main_job_ids]
 
         return result
 


Test patch:
diff --git a/tests/test_db_inquirer.py b/tests/test_db_inquirer.py
index 8742387..209b145 100644
--- a/tests/test_db_inquirer.py
+++ b/tests/test_db_inquirer.py
@@ -613,7 +613,7 @@ def test_sacct_get_db_output_user_state(sacct, mocker):
     mock_date.side_effect = datetime.date
     mocker.patch("reportseff.db_inquirer.datetime.date", mock_date)
     with pytest.raises(Exception, match="Error running sacct!"):
-        sacct.get_db_output("c1 c2".split(), "j1 j2 j3".split())
+        sacct.get_db_output("JobID c2".split(), "j1 j2 j3".split())
 
     mock_sacct = mocker.MagicMock()
     mock_sacct.returncode = 0
@@ -625,14 +625,14 @@ def test_sacct_get_db_output_user_state(sacct, mocker):
     )
     sacct.set_user("user")
     sacct.set_state("R")
-    result = sacct.get_db_output("c1 c2 State".split(), {})
+    result = sacct.get_db_output("JobID c2 State".split(), {})
     assert result == [
-        {"c1": "c1j1", "c2": "c2j1", "State": "RUNNING"},
-        {"c1": "c1j2", "c2": "c2j2", "State": "RUNNING"},
+        {"JobID": "c1j1", "c2": "c2j1", "State": "RUNNING"},
+        {"JobID": "c1j2", "c2": "c2j2", "State": "RUNNING"},
     ]
     mock_sub.assert_called_once_with(
         args=(
-            "sacct -P -n --delimiter=^|^ --format=c1,c2,State"
+            "sacct -P -n --delimiter=^|^ --format=JobID,c2,State"
             " --user=user --starttime=011318"
         ).split(),
         stdout=mocker.ANY,
@@ -644,7 +644,7 @@ def test_sacct_get_db_output_user_state(sacct, mocker):
 
     # debug is not affected by state
     debug = []
-    sacct.get_db_output("c1 c2 State".split(), "j1 j2 j3".split(), debug.append)
+    sacct.get_db_output("JobID c2 State".split(), "j1 j2 j3".split(), debug.append)
     assert debug[0] == (
         "c1j1^|^c2j1^|^RUNNING\nc1j2^|^c2j2^|^RUNNING\nc1j3^|^c2j3^|^COMPLETED\n"
     )
diff --git a/tests/test_reportseff.py b/tests/test_reportseff.py
index c070712..f0b6b73 100644
--- a/tests/test_reportseff.py
+++ b/tests/test_reportseff.py
@@ -892,3 +892,46 @@ def test_extra_args(mocker):
     )
 
     assert result.exit_code == 0
+
+
+@pytest.mark.usefixtures("_mock_inquirer")
+def test_issue_58(mocker):
+    """Incorrect memory usage when filtering by state."""
+    mocker.patch("reportseff.console.which", return_value=True)
+    runner = CliRunner()
+    sub_result = mocker.MagicMock()
+    sub_result.returncode = 0
+    sub_result.stdout = """
+^|^2^|^01:00:29^|^1234^|^1234^|^^|^1^|^^|^40G^|^TIMEOUT^|^01:00:00^|^01:01:06
+^|^2^|^01:00:33^|^1234.batch^|^1234.batch^|^19855760K^|^1^|^1^|^^|^CANCELLED^|^^|^01:01:06
+"""
+    mocker.patch("reportseff.db_inquirer.subprocess.run", return_value=sub_result)
+    result = runner.invoke(console.main, "--no-color 1234".split())
+
+    assert result.exit_code == 0
+    # remove header
+    output = result.output.split("\n")[1:-1]
+    assert output[0].split() == [
+        "1234",
+        "TIMEOUT",
+        "01:00:29",
+        "100.8%",
+        "50.5%",
+        "47.3%",
+    ]
+    assert len(output) == 1
+
+    result = runner.invoke(console.main, "--state TIMEOUT --no-color 1234".split())
+
+    assert result.exit_code == 0
+    # remove header
+    output = result.output.split("\n")[1:-1]
+    assert output[0].split() == [
+        "1234",
+        "TIMEOUT",
+        "01:00:29",
+        "100.8%",
+        "50.5%",
+        "47.3%",
+    ]
+    assert len(output) == 1