1  """
  2  Console module tests
  3  """
  4  
  5  import contextlib
  6  import io
  7  import os
  8  import tempfile
  9  import unittest
 10  
 11  from txtai.console import Console
 12  from txtai.embeddings import Embeddings
 13  
 14  APPLICATION = """
 15  path: %s
 16  
 17  workflow:
 18    test:
 19       tasks:
 20         - task: console
 21  """
 22  
 23  
 24  class TestConsole(unittest.TestCase):
 25      """
 26      Console tests.
 27      """
 28  
 29      @classmethod
 30      def setUpClass(cls):
 31          """
 32          Initialize test data.
 33          """
 34  
 35          cls.data = [
 36              "US tops 5 million confirmed virus cases",
 37              "Canada's last fully intact ice shelf has suddenly collapsed, forming a Manhattan-sized iceberg",
 38              "Beijing mobilises invasion craft along coast as Taiwan tensions escalate",
 39              "The National Park Service warns against sacrificing slower friends in a bear attack",
 40              "Maine man wins $1M from $25 lottery ticket",
 41              "Make huge profits without work, earn up to $100,000 a day",
 42          ]
 43  
 44          # Create embeddings model, backed by sentence-transformers & transformers
 45          cls.embeddings = Embeddings({"path": "sentence-transformers/nli-mpnet-base-v2", "content": True})
 46  
 47          # Create an index for the list of text
 48          cls.embeddings.index([(uid, text, None) for uid, text in enumerate(cls.data)])
 49  
 50          # Create app paths
 51          cls.apppath = os.path.join(tempfile.gettempdir(), "console.yml")
 52          cls.embedpath = os.path.join(tempfile.gettempdir(), "embeddings.console")
 53  
 54          # Create app.yml
 55          with open(cls.apppath, "w", encoding="utf-8") as out:
 56              out.write(APPLICATION % cls.embedpath)
 57  
 58          # Save index as uncompressed and compressed
 59          cls.embeddings.save(cls.embedpath)
 60          cls.embeddings.save(f"{cls.embedpath}.tar.gz")
 61  
 62          # Create console
 63          cls.console = Console(cls.embedpath)
 64  
 65      def testApplication(self):
 66          """
 67          Test application
 68          """
 69  
 70          self.assertNotIn("Traceback", self.command(f".load {self.apppath}"))
 71          self.assertIn("1", self.command(".limit 1"))
 72          self.assertIn("Maine man wins", self.command("feel good story"))
 73  
 74      def testConfig(self):
 75          """
 76          Test .config command
 77          """
 78  
 79          self.assertIn("tasks", self.command(".config"))
 80  
 81      def testEmbeddings(self):
 82          """
 83          Test embeddings index
 84          """
 85  
 86          self.assertNotIn("Traceback", self.command(f".load {self.embedpath}.tar.gz"))
 87          self.assertNotIn("Traceback", self.command(f".load {self.embedpath}"))
 88          self.assertIn("1", self.command(".limit 1"))
 89          self.assertIn("Maine man wins", self.command("feel good story"))
 90  
 91      def testEmbeddingsNoDatabase(self):
 92          """
 93          Test embeddings with no database/content
 94          """
 95  
 96          console = Console()
 97  
 98          # Create embeddings model, backed by sentence-transformers & transformers
 99          embeddings = Embeddings({"path": "sentence-transformers/nli-mpnet-base-v2"})
100  
101          # Create an index for the list of text
102          embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])
103  
104          # Set embeddings on console
105          console.app = embeddings
106          self.assertIn("4", self.command("feel good story", console))
107  
108      def testEmpty(self):
109          """
110          Test empty console instance
111          """
112  
113          console = Console()
114          self.assertIn("AttributeError", self.command("search", console))
115  
116      def testHighlight(self):
117          """
118          Test .highlight command
119          """
120  
121          self.assertIn("highlight", self.command(".highlight"))
122          self.assertIn("wins", self.command("feel good story"))
123          self.assertIn("Taiwan", self.command("asia"))
124  
125      def testPreloop(self):
126          """
127          Test preloop
128          """
129  
130          self.assertIn("txtai console", self.preloop())
131  
132      def testWorkflow(self):
133          """
134          Test .workflow command
135          """
136  
137          self.command(f".load {self.apppath}")
138          self.assertIn("echo", self.command(".workflow test echo"))
139  
140      def command(self, command, console=None):
141          """
142          Runs a console command.
143  
144          Args:
145              command: command to run
146              console: console instance, defaults to self.console
147  
148          Returns:
149              command output
150          """
151  
152          # Run info
153          output = io.StringIO()
154          with contextlib.redirect_stdout(output):
155              if not console:
156                  console = self.console
157  
158              console.onecmd(command)
159  
160          return output.getvalue()
161  
162      def preloop(self):
163          """
164          Runs console.preloop and redirects stdout.
165  
166          Returns:
167              preloop output
168          """
169  
170          # Run info
171          output = io.StringIO()
172          with contextlib.redirect_stdout(output):
173              self.console.preloop()
174  
175          return output.getvalue()