# -*- coding: utf-8 -*-

import os

import pybel
from pybel.constants import *
from pybel_tools.mutation import prune_central_dogma, infer_missing_inverse_edge, infer_central_dogma
from tests.constants import ManagerMixin


class TestProcessing(ManagerMixin):
    def setUp(self):
        super(TestProcessing, self).setUp()

        if 'PYBEL_BASE' in os.environ:
            test_bel_simple_path = os.path.join(os.environ['PYBEL_BASE'], 'tests', 'bel', 'test_bel.bel')
            self.graph = pybel.from_path(test_bel_simple_path, manager=self.manager)
        else:
            test_bel_simple_url = 'https://raw.githubusercontent.com/pybel/pybel/develop/tests/bel/test_bel.bel'
            self.graph = pybel.from_url(test_bel_simple_url, manager=self.manager)

        infer_central_dogma(self.graph)

        n1 = GENE, 'HGNC', 'AKT1'
        n2 = RNA, 'HGNC', 'EGFR'

        n3 = GENE, 'HGNC', 'DUMMY1'
        self.graph.add_simple_node(*n3)

        n4 = GENE, 'HGNC', 'DUMMY2'
        self.graph.add_simple_node(*n4)

        self.graph.add_edge(n1, n3, **{RELATION: INCREASES})
        self.graph.add_edge(n2, n4, **{RELATION: INCREASES})

    def test_prune(self):
        self.assertEqual(14, self.graph.number_of_nodes())
        self.assertEqual(16, self.graph.number_of_edges())

        prune_central_dogma(self.graph)
        self.assertEqual(9, self.graph.number_of_nodes())

    def test_infer(self):
        self.assertEqual(14, self.graph.number_of_nodes())
        self.assertEqual(16, self.graph.number_of_edges())

        infer_missing_inverse_edge(self.graph, TRANSLATED_TO)
        self.assertEqual(20, self.graph.number_of_edges())

# -*- coding: utf-8 -*-

import unittest

from pybel.dsl import gene, rna
from pybel.examples import sialic_acid_graph
from pybel_tools.mutation import infer_central_dogma, prune_central_dogma

trem2_gene = gene(namespace='HGNC', name='TREM2')
trem2_rna = rna(namespace='HGNC', name='TREM2')


class TestProcessing(unittest.TestCase):
    def test_infer_on_sialic_acid_example(self):
        sgc = sialic_acid_graph.copy()

        self.assertFalse(sgc.has_node_with_data(trem2_gene))
        self.assertFalse(sgc.has_node_with_data(trem2_rna))

        infer_central_dogma(sgc)

        self.assertTrue(sgc.has_node_with_data(trem2_gene))
        self.assertTrue(sgc.has_node_with_data(trem2_rna))

        prune_central_dogma(sgc)

        self.assertFalse(sgc.has_node_with_data(trem2_gene))
        self.assertFalse(sgc.has_node_with_data(trem2_rna))

Fix testing