-
Notifications
You must be signed in to change notification settings - Fork 208
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Merge pull request #12 from lyft/tfeng_update_user_model
Add User node with different user attributes
- Loading branch information
Showing
7 changed files
with
240 additions
and
26 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,161 @@ | ||
| from typing import Union, Dict, Any # noqa: F401 | ||
|
|
||
| from databuilder.models.neo4j_csv_serde import Neo4jCsvSerializable, NODE_KEY, \ | ||
| NODE_LABEL, RELATION_START_KEY, RELATION_START_LABEL, RELATION_END_KEY, \ | ||
| RELATION_END_LABEL, RELATION_TYPE, RELATION_REVERSE_TYPE | ||
|
|
||
|
|
||
| class User(Neo4jCsvSerializable): | ||
| # type: (...) -> None | ||
| """ | ||
| User model. This model doesn't define any relationship. | ||
| """ | ||
| USER_NODE_LABEL = 'User' | ||
| USER_NODE_KEY_FORMAT = '{email}' | ||
| USER_NODE_EMAIL = 'email' | ||
| USER_NODE_FIRST_NAME = 'first_name' | ||
| USER_NODE_LAST_NAME = 'last_name' | ||
| USER_NODE_FULL_NAME = 'full_name' | ||
| USER_NODE_GITHUB_NAME = 'github_username' | ||
| USER_NODE_TEAM = 'team_name' | ||
| USER_NODE_EMPLOYEE_TYPE = 'employee_type' | ||
| USER_NODE_MANAGER_EMAIL = 'manager_email' | ||
| USER_NODE_SLACK_ID = 'slack_id' | ||
| USER_NODE_IS_ACTIVE = 'is_active' | ||
| USER_NODE_UPDATED_AT = 'updated_at' | ||
|
|
||
| USER_MANAGER_RELATION_TYPE = 'MANAGE_BY' | ||
| MANAGER_USER_RELATION_TYPE = 'MANAGE' | ||
|
|
||
| def __init__(self, | ||
| email, # type: str | ||
| first_name='', # type: str | ||
| last_name='', # type: str | ||
| name='', # type: str | ||
| github_username='', # type: str | ||
| team_name='', # type: str | ||
| employee_type='', # type: str | ||
| manager_email='', # type: str | ||
| slack_id='', # type: str | ||
| is_active=True, # type: bool | ||
| updated_at=0, # type: int | ||
| ): | ||
| # type: (...) -> None | ||
| """ | ||
| This class models user node for Amundsen people. | ||
| :param first_name: | ||
| :param last_name: | ||
| :param name: | ||
| :param email: | ||
| :param github_username: | ||
| :param team_name: | ||
| :param employee_type: | ||
| :param manager_email: | ||
| :param is_active: | ||
| :param updated_at: everytime we update the node, we will push the timestamp. | ||
| then we will have a cron job to update the ex-employee nodes based on | ||
| the case if this timestamp hasn't been updated for two weeks. | ||
| """ | ||
| self.first_name = first_name | ||
| self.last_name = last_name | ||
| self.name = name | ||
| self.email = email | ||
| self.github_username = github_username | ||
| # todo: team will be a separate node once Amundsen People supports team | ||
| self.team_name = team_name | ||
| self.manager_email = manager_email | ||
| self.employee_type = employee_type | ||
| # this attr not available in team service, either update team service, update with FE | ||
| self.slack_id = slack_id | ||
| self.is_active = is_active | ||
| self.updated_at = updated_at | ||
|
|
||
| self._node_iter = iter(self.create_nodes()) | ||
| self._rel_iter = iter(self.create_relation()) | ||
|
|
||
| def create_next_node(self): | ||
| # type: (...) -> Union[Dict[str, Any], None] | ||
| # return the string representation of the data | ||
| try: | ||
| return next(self._node_iter) | ||
| except StopIteration: | ||
| return None | ||
|
|
||
| def create_next_relation(self): | ||
| # type: () -> Union[Dict[str, Any], None] | ||
| """ | ||
| :return: | ||
| """ | ||
| try: | ||
| return next(self._rel_iter) | ||
| except StopIteration: | ||
| return None | ||
|
|
||
| @classmethod | ||
| def get_user_model_key(cls, | ||
| email=None): | ||
| # type: (...) -> str | ||
| if not email: | ||
| return '' | ||
| return User.USER_NODE_KEY_FORMAT.format(email=email) | ||
|
|
||
| def create_nodes(self): | ||
| # type: () -> List[Dict[str, Any]] | ||
| """ | ||
| Create a list of Neo4j node records | ||
| :return: | ||
| """ | ||
| result_node = { | ||
| NODE_KEY: User.get_user_model_key(email=self.email), | ||
| NODE_LABEL: User.USER_NODE_LABEL, | ||
| User.USER_NODE_EMAIL: self.email, | ||
| User.USER_NODE_IS_ACTIVE: self.is_active, | ||
| } | ||
|
|
||
| if self.first_name: | ||
| result_node[User.USER_NODE_FIRST_NAME] = self.first_name | ||
| if self.last_name: | ||
| result_node[User.USER_NODE_LAST_NAME] = self.last_name | ||
| if self.name: | ||
| result_node[User.USER_NODE_FULL_NAME] = self.name | ||
| if self.github_username: | ||
| result_node[User.USER_NODE_GITHUB_NAME] = self.github_username | ||
| if self.team_name: | ||
| result_node[User.USER_NODE_TEAM] = self.team_name | ||
| if self.employee_type: | ||
| result_node[User.USER_NODE_EMPLOYEE_TYPE] = self.employee_type | ||
| if self.slack_id: | ||
| result_node[User.USER_NODE_SLACK_ID] = self.slack_id | ||
| if self.updated_at: | ||
| result_node[User.USER_NODE_UPDATED_AT] = self.updated_at | ||
|
|
||
| return [result_node] | ||
|
|
||
| def create_relation(self): | ||
| # type: () -> List[Dict[str, Any]] | ||
| if self.manager_email: | ||
| # only create the relation if the manager exists | ||
| return [{ | ||
| RELATION_START_KEY: User.get_user_model_key(email=self.email), | ||
| RELATION_START_LABEL: User.USER_NODE_LABEL, | ||
| RELATION_END_KEY: self.get_user_model_key(email=self.manager_email), | ||
| RELATION_END_LABEL: User.USER_NODE_LABEL, | ||
| RELATION_TYPE: User.USER_MANAGER_RELATION_TYPE, | ||
| RELATION_REVERSE_TYPE: User.MANAGER_USER_RELATION_TYPE | ||
| }] | ||
| return [] | ||
|
|
||
| def __repr__(self): | ||
| # type: () -> str | ||
| return 'User({!r}, {!r}, {!r}, {!r}, {!r}, ' \ | ||
| '{!r}, {!r}, {!r}, {!r}, {!r}, {!r},)'.format(self.first_name, | ||
| self.last_name, | ||
| self.name, | ||
| self.email, | ||
| self.github_username, | ||
| self.team_name, | ||
| self.slack_id, | ||
| self.manager_email, | ||
| self.employee_type, | ||
| self.is_active) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -22,8 +22,14 @@ def test_serialize(self): | |
| actual.append(node_row) | ||
| node_row = table_col_usage.next_node() | ||
|
|
||
| expected = [{'email': '[email protected]', 'KEY': '[email protected]', 'LABEL': 'User'}, | ||
| {'email': '[email protected]', 'KEY': '[email protected]', 'LABEL': 'User'}] | ||
| expected = [{'is_active': True, | ||
| 'LABEL': 'User', | ||
| 'KEY': '[email protected]', | ||
| 'email': '[email protected]'}, | ||
| {'is_active': True, | ||
| 'LABEL': 'User', | ||
| 'KEY': '[email protected]', | ||
| 'email': '[email protected]'}] | ||
| self.assertEqual(expected, actual) | ||
|
|
||
| rel_row = table_col_usage.next_relation() | ||
|
|
||
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,53 @@ | ||
| import unittest | ||
|
|
||
| from databuilder.models.neo4j_csv_serde import RELATION_START_KEY, RELATION_START_LABEL, RELATION_END_KEY, \ | ||
| RELATION_END_LABEL, RELATION_TYPE, RELATION_REVERSE_TYPE | ||
|
|
||
| from databuilder.models.user import User | ||
|
|
||
|
|
||
| class TestUser(unittest.TestCase): | ||
|
|
||
| def setUp(self): | ||
| # type: () -> None | ||
| super(TestUser, self).setUp() | ||
| self.user = User(first_name='test_first', | ||
| last_name='test_last', | ||
| name='test_first test_last', | ||
| email='[email protected]', | ||
| github_username='github_test', | ||
| team_name='test_team', | ||
| employee_type='FTE', | ||
| manager_email='[email protected]', | ||
| slack_id='slack', | ||
| is_active=True, | ||
| updated_at=1) | ||
|
|
||
| def test_get_user_model_key(self): | ||
| # type: () -> None | ||
| user_email = User.get_user_model_key(email=self.user.email) | ||
| self.assertEquals(user_email, '{email}'.format(email='[email protected]')) | ||
|
|
||
| def test_create_nodes(self): | ||
| # type: () -> None | ||
| nodes = self.user.create_nodes() | ||
| self.assertEquals(len(nodes), 1) | ||
|
|
||
| def test_create_relation(self): | ||
| # type: () -> None | ||
| relations = self.user.create_relation() | ||
| self.assertEquals(len(relations), 1) | ||
|
|
||
| start_key = '{email}'.format(email='[email protected]') | ||
| end_key = '{email}'.format(email='[email protected]') | ||
|
|
||
| relation = { | ||
| RELATION_START_KEY: start_key, | ||
| RELATION_START_LABEL: User.USER_NODE_LABEL, | ||
| RELATION_END_KEY: end_key, | ||
| RELATION_END_LABEL: User.USER_NODE_LABEL, | ||
| RELATION_TYPE: User.USER_MANAGER_RELATION_TYPE, | ||
| RELATION_REVERSE_TYPE: User.MANAGER_USER_RELATION_TYPE | ||
| } | ||
|
|
||
| self.assertTrue(relation in relations) |