from decagon_pytorch.data import Data
from decagon_pytorch.decode import DEDICOMDecoder


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')
    d.set_decoder_type(1, 1, DEDICOMDecoder)
    print(d)