friends.py

#!/usr/bin/env py.test

from collections import deque, defaultdict

class Friends:
    def __init__(self):
        self.graph = defaultdict(set)

    def make_friend(self, name1, name2):
        self.graph[name1].add(name2)
        self.graph[name2].add(name1)

    def unmake_friend(self, name1, name2):
        self.graph[name1].discard(name2)
        self.graph[name2].discard(name1)

    def get_direct_friends(self, name):
        return self.graph[name]

    def get_indirect_friends(self, name):
        queue = deque(self.graph[name])
        found_friends = set([name])
        while queue:
            friend = queue.popleft()
            if friend not in found_friends:
                yield friend
                found_friends.add(friend)
                queue.extend(self.graph[friend])




import pytest

@pytest.fixture
def graph():
    return Friends()

@pytest.fixture
def graph1(graph):
    graph.make_friend("Aaron", "Bella")
    graph.make_friend("Cindy", "Bella")
    graph.make_friend("David", "Bella")
    graph.make_friend("David", "Elizabeth")
    graph.make_friend("Cindy", "Frank")
    return graph

def test_friend_pair(graph):
    assert set() == graph.get_direct_friends("I don't exist")
    graph.make_friend("Aaron", "Bella")
    assert "Bella" in graph.get_direct_friends("Aaron")
    assert "Aaron" in graph.get_direct_friends("Bella")

def test_friend_3tuple(graph):
    graph.make_friend("Aaron", "Bella")
    graph.make_friend("Bella", "Frank")
    assert "Bella" in graph.get_direct_friends("Aaron")
    assert "Bella" in graph.get_direct_friends("Frank")
    assert {"Aaron", "Frank"} == graph.get_direct_friends("Bella")

def test_direct_friends(graph1):
    assert {'Aaron', 'Cindy', 'David'} == graph1.get_direct_friends('Bella')

def test_indirect_friends(graph1):
    expected_frank_friends = {'Aaron', 'Bella', 'Cindy', 'David', 'Elizabeth'}
    assert expected_frank_friends == set(graph1.get_indirect_friends('Frank'))

def test_unmake_friend(graph1):
    graph1.unmake_friend('Bella', 'David')
    assert {'Bella', 'Cindy', 'Aaron'} == set(graph1.get_indirect_friends('Frank'))
    assert {'David'} == set(graph1.get_indirect_friends('Elizabeth'))

def test_no_duplicates(graph1):
    graph1.make_friend('Cindy', 'David')
    frank_friends = list(graph1.get_indirect_friends('Frank'))
    assert len(frank_friends) == len(set(frank_friends))

def test_unmake_friend2(graph1):
    graph1.make_friend('Cindy', 'David')
    graph1.unmake_friend('Bella', 'David')
    expected_frank_friends = {'Aaron', 'Bella', 'Cindy', 'David', 'Elizabeth'}
    assert expected_frank_friends == set(graph1.get_indirect_friends('Frank'))
    graph1.unmake_friend('Bella', 'Cindy')
    assert {'Bella'} == set(graph1.get_indirect_friends('Aaron'))
    assert {'Cindy', 'David', 'Frank'} == set(graph1.get_indirect_friends('Elizabeth'))

requirements.txt

py==1.4.23
pytest==2.6.1