from decagon_pytorch.data import Data


def test_data():
    d = Data()
    d.add_node_type('Gene')
    d.add_node_type('Drug')
    d.add_relation(1, 0, None, 'Target')
    d.add_relation(0, 0, None, 'Interaction')
    d.add_relation(1, 1, None, 'Side Effect: Nausea')
    d.add_relation(1, 1, None, 'Side Effect: Infertility')
    d.add_relation(1, 1, None, 'Side Effect: Death')
    print(d)