Coverage for tests/test_sentiment.py: 97%

1 import pytest

2 import mall

3 import polars as pl

4 import pyarrow

5

6 import shutil

7 import os

8

9 if os._exists("_test_cache"):

10 shutil.rmtree("_test_cache", ignore_errors=True)

11

12

13 def test_sentiment_simple():

14 data = mall.MallData

15 reviews = data.reviews

16 reviews.llm.use("test", "echo", _cache="_test_cache")

17 x = reviews.llm.sentiment("review")

18 assert (

19 x.select("sentiment").to_pandas().to_string()

20 == " sentiment\n0 None\n1 None\n2 None"

21 )

22

23

24 def sim_sentiment():

25 df = pl.DataFrame(dict(x=["positive", "negative", "neutral", "not-real"]))

26 df.llm.use("test", "echo", _cache="_test_cache")

27 return df

28

29

30 def test_sentiment_valid():

31 x = sim_sentiment()

32 x = x.llm.sentiment("x")

33 assert (

34 x.select("sentiment").to_pandas().to_string()

35 == " sentiment\n0 positive\n1 negative\n2 neutral\n3 None"

36 )

37

38

39 def test_sentiment_valid2():

40 x = sim_sentiment()

41 x = x.llm.sentiment("x", ["positive", "negative"])

42 assert (

43 x.select("sentiment").to_pandas().to_string()

44 == " sentiment\n0 positive\n1 negative\n2 None\n3 None"

45 )

46

47

48 def test_sentiment_prompt():

49 df = pl.DataFrame(dict(x="x"))

50 df.llm.use("test", "content", _cache="_test_cache")

51 x = df.llm.sentiment("x")

52 assert (

53 x["sentiment"][0]

54 == "You are a helpful sentiment engine. Return only one of the following answers: positive, negative, neutral . No capitalization. No explanations. The answer is based on the following text:\n{}"

55 )