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

"""Tests for node selection functions."""

import unittest

from pybel import BELGraph
from pybel.dsl import Protein
from pybel.struct.filters import get_nodes_by_namespace


class TestNodeSelection(unittest.TestCase):
    """Tests for node selection functions."""

    def test_get_node_by_namespace(self):
        """Test getting nodes with a given namespace."""
        g = BELGraph()
        a = Protein(namespace="N1", name="a")
        b = Protein(namespace="N1", name="b")
        c = Protein(namespace="N2", name="c")
        d = Protein(namespace="N3", name="d")
        g.add_node_from_data(a)
        g.add_node_from_data(b)
        g.add_node_from_data(c)
        g.add_node_from_data(d)

        nodes = set(get_nodes_by_namespace(g, "N1"))

        self.assertIn(a, nodes)
        self.assertIn(b, nodes)
        self.assertNotIn(c, nodes)
        self.assertNotIn(d, nodes)

        nodes = set(get_nodes_by_namespace(g, ("N1", "N2")))

        self.assertIn(a, nodes)
        self.assertIn(b, nodes)
        self.assertIn(c, nodes)
        self.assertNotIn(d, nodes)
