Coverage for mindsdb / integrations / handlers / unify_handler / tests / test_unify_handler.py: 0%

1import os 

2import pytest 

3 

4from mindsdb_sql_parser import parse_sql 

5 

6from tests.unit.executor_test_base import BaseExecutorTest 

7 

8 

9class TestUnify(BaseExecutorTest): 

10 

11 def run_sql(self, sql): 

12 ret = self.command_executor.execute_command(parse_sql(sql)) 

13 assert ret.error_code is None 

14 if ret.data is not None: 

15 return ret.data.to_df() 

16 

17 """ 

18 Integration tests for the Unify handler. 

19 """ 

20 @pytest.fixture(autouse=True, scope="function") 

21 def setup_method(self): 

22 """ 

23 Setup test environment, creating a project 

24 """ 

25 super().setup_method() 

26 self.run_sql("create database proj") 

27 

28 @pytest.mark.skipif(os.environ.get('UNIFY_API_KEY') is None, reason='Missing API key!') 

29 def test_unify_correct_flow(self): 

30 self.run_sql( 

31 f""" 

32 CREATE ML_ENGINE unify_engine 

33 FROM unify 

34 USING 

35 unify_api_key = '{os.environ.get('UNIFY_API_KEY')}'; 

36 """ 

37 ) 

38 

39 self.run_sql( 

40 """ 

41 CREATE MODEL proj.test_unify_correct_flow 

42 PREDICT output 

43 USING 

44 engine='unify_engine', 

45 model = 'llama-3-8b-chat', 

46 provider = 'together-ai', 

47 question_column = 'text'; 

48 """ 

49 ) 

50 self.wait_predictor("proj", "test_unify_correct_flow") 

51 

52 result_df = self.run_sql( 

53 """ 

54 SELECT text, output 

55 FROM proj.test_unify_correct_flow 

56 WHERE text = 'Hello'; 

57 """ 

58 ) 

59 assert "hello" in result_df["output"].iloc[0].lower()