Given an org chart, debug it and then print it in JSON format

org_chart.py

""""
Imagine you've been asked to review, and possibly refactor, the below code (of course, for the purpose of this interview, the code is intentionally poorly written). Think of your interviewer as your pair programming partner. You are expected to lead the exercise, but they are there to help you and provide guidance where appropriate.

Start by understanding and explaining what the code does. Feel free to make small changes as you go. There are also tests at the end that may help.

Next, try to get any failing tests to pass.

From there, come up with a plan for how you might improve the code. There may be both code and performance improvements.

We will block these segments by time, so don't be alarmed if we need to move to the next part.
"""

import json
import unittest

class Employee:
  def __init__(self, name):
    self.name = name
    self.manager = None
    self.reports = []

  def __repr__(self):
    return self.name

def build_employee_dummy_data():
  ada = Employee('Ada')
  booth = Employee('Booth')
  conway = Employee('Conway')
  dahl = Employee('Dahl')
  evi = Employee('Evi')
  fran = Employee('Fran')

  ada.manager = booth
  dahl.manager = booth
  conway.manager = ada
  evi.manager = dahl
  fran.manager = dahl

  booth.reports = [ada, dahl]
  ada.reports = [conway]
  dahl.reports = [evi, fran]

  return [fran, ada, booth, conway, dahl, evi]

def print_org_chart(employees):
  q = []
  for employee in employees:
    if not employee.manager:
      q.append(employee)

  employee = q.pop(0)
    
  output = ""
  while employee:
    # Handle all employees that have reports and print those relationships breadth-first
    if employee.reports is not None and len(employee.reports) > 0:
      output += employee.name + " -> "
     
      for emp in employee.reports:
        output += emp.name + " "
        q.append(emp)

      output += "\n"
    # We want to print lone ranger employees but not low-level individual contributor reports
    elif employee.manager is None:
      output += employee.name + "\n"

    if q:
      employee = q.pop(0)
    else:
      employee = None
    
  return output

# {
#   "booth": {
#     "ada": {
#       "conway": {}
#     },
#     "dahl": {
#       "evi": {},
#       "fran": {}
#     }
#   }
# }

def print_json_org_chart(employees):
  direct_reports_of = {}
  for employee in employees:
    direct_reports_of[employee.name] = employee.reports

  # Recursively determine the org chart with this employee at the top
  def get_employee_subtree(employee_name):
    res = {}

    for report in direct_reports_of[employee_name]:
        res[str(report)] = get_employee_subtree(report.name)

    return res

  # Iterate through all the top-level managers one-by-one and generate their org tree
  res = {}
  for employee in employees:
    if employee.manager is None:
      res[employee.name] = get_employee_subtree(employee.name)

  return json.dumps(res)



### Unit Tests ###

class TestPrintOrgChart(unittest.TestCase):

  def test_every_employee_has_a_manager_or_direct_reports(self):
    expected_output = ("Booth -> Ada Dahl \n"
                       "Ada -> Conway \n"
                       "Dahl -> Evi Fran \n")
    employees = build_employee_dummy_data()
    self.assertEqual(print_org_chart(employees), expected_output)


  def test_at_least_one_employee_with_no_manager_or_direct_reports(self):
    employees = [Employee("Mr Magoo"), *build_employee_dummy_data()]
    expected_output = ("Mr Magoo\n"
                       "Booth -> Ada Dahl \n"
                       "Ada -> Conway \n"
                       "Dahl -> Evi Fran \n")
    self.assertEqual(print_org_chart(employees), expected_output)


  def test_json_every_employee_has_a_manager_or_direct_reports(self):
    expected_output = ("{\"Booth\": {\"Ada\": {\"Conway\": {}}, \"Dahl\": {\"Evi\": {}, \"Fran\": {}}}}")
    employees = build_employee_dummy_data()
    self.assertEqual(print_json_org_chart(employees), expected_output)


  def test_json_at_least_one_employee_with_no_manager_or_direct_reports(self):
    employees = [Employee("Mr Magoo"), *build_employee_dummy_data()]
    expected_output = ("{\"Mr Magoo\": {}, \"Booth\": {\"Ada\": {\"Conway\": {}}, \"Dahl\": {\"Evi\": {}, \"Fran\": {}}}}")
    self.assertEqual(print_json_org_chart(employees), expected_output)

unittest.main(exit=False)