refactor

Author: devadathanmb

pgcli/main.py

@@ -1248,6 +1248,20 @@ def get_completions(self, text, cursor_positition):
         with self._completer_lock:
             return self.completer.get_completions(Document(text=text, cursor_position=cursor_positition), None)
 
+    def _get_transaction_status_char(self) -> str:
+        """Return transaction status character for prompt.
+
+        Following psql convention:
+        - '*' when in a valid transaction (ACTIVE or INTRANS)
+        - '!' when in a failed transaction (INERROR)
+        - '' when idle
+        """
+        if self.pgexecute.failed_transaction():
+            return "!"
+        if self.pgexecute.valid_transaction():
+            return "*"
+        return ""
+
     def get_prompt(self, string):
         # should be before replacing \\d
         string = string.replace("\\dsn_alias", self.dsn_alias or "")
@@ -1263,7 +1277,7 @@ def get_prompt(self, string):
         string = string.replace("\\i", str(self.pgexecute.pid) or "(none)")
         string = string.replace("\\#", "#" if self.pgexecute.superuser else ">")
         string = string.replace("\\n", "\n")
-        string = string.replace("\\x", self.pgexecute.transaction_status())
+        string = string.replace("\\x", self._get_transaction_status_char())
         return string
 
     def get_last_query(self):

pgcli/pgexecute.py

@@ -298,20 +298,6 @@ def valid_transaction(self):
         status = self.conn.info.transaction_status
         return status == psycopg.pq.TransactionStatus.ACTIVE or status == psycopg.pq.TransactionStatus.INTRANS
 
-    def transaction_status(self):
-        """Return transaction status character for prompt.
-
-        Following psql convention:
-        - '*' when in a valid transaction (ACTIVE or INTRANS)
-        - '!' when in a failed transaction (INERROR)
-        - '' when idle
-        """
-        if self.failed_transaction():
-            return "!"
-        elif self.valid_transaction():
-            return "*"
-        return ""
-
     def run(
         self,
         statement,

tests/test_main.py

@@ -578,14 +578,14 @@ def test_duration_in_words(duration_in_seconds, words):
 
 
 @pytest.mark.parametrize(
-    "transaction_status,expected",
+    "is_failed,is_valid,expected",
     [
-        ("*", "*testuser"),
-        ("!", "!testuser"),
-        ("", "testuser"),
+        (False, True, "*testuser"),  # valid transaction → "*"
+        (True, False, "!testuser"),  # failed transaction → "!"
+        (False, False, "testuser"),  # idle → ""
     ],
 )
-def test_get_prompt_with_transaction_status(transaction_status, expected):
+def test_get_prompt_with_transaction_status(is_failed, is_valid, expected):
     """Test that \\x prompt variable shows transaction status."""
     cli = PGCli()
     cli.pgexecute = mock.MagicMock()
@@ -597,10 +597,10 @@ def test_get_prompt_with_transaction_status(transaction_status, expected):
     cli.pgexecute.pid = 12345
     cli.pgexecute.superuser = False
 
-    with mock.patch.object(
-        cli.pgexecute, "transaction_status", return_value=transaction_status
-    ):
-        result = cli.get_prompt("\\x\\u")
+    cli.pgexecute.failed_transaction.return_value = is_failed
+    cli.pgexecute.valid_transaction.return_value = is_valid
+
+    result = cli.get_prompt("\\x\\u")
     assert result == expected
 
 
@@ -616,10 +616,10 @@ def test_get_prompt_transaction_status_in_full_prompt():
     cli.pgexecute.pid = 12345
     cli.pgexecute.superuser = False
 
-    with mock.patch.object(
-        cli.pgexecute, "transaction_status", return_value="*"
-    ):
-        result = cli.get_prompt("\\x\\u@\\h:\\d> ")
+    cli.pgexecute.failed_transaction.return_value = False
+    cli.pgexecute.valid_transaction.return_value = True
+
+    result = cli.get_prompt("\\x\\u@\\h:\\d> ")
     assert result == "*user@db.example.com:mydb> "