import polars as pl
from smithy import smith_set

def test_condorcet():
    df = pl.DataFrame({
                          'A': [1, 1, 2, 1],
                          'B': [2, 2, 1, 2],
                          'C': [3, 3, 3, 3],
                      })
    assert smith_set(df) == ['A']

def test_rockpprscrcycle():
    df = pl.DataFrame({
                             'A': [1, 2, 3],
                             'B': [2, 3, 1],
                             'C': [3, 1, 2],
                         })
    assert smith_set(df) == ['A', 'B', 'C']

def test_abpair():
    df = pl.DataFrame({
        "A": [1, 2, 1, 3],
        "B": [2, 1, 3, 1],
        "C": [3, 3, 2, 2]
        })
    assert smith_set(df) == ['A', 'B']

def test_fourcycle():
    df = pl.DataFrame({
        "A": [1,2,3,4],
        "B": [2,3,4,1],
        "C": [3,4,1,2],
        "D": [4,1,2,3],
    })
    assert smith_set(df) == ['A', 'B', 'C', 'D']