Coverage for compiler_admin / commands / ls / users.py: 100%

1import click 

2import pandas as pd 

3 

4from compiler_admin import FORMATS, Format 

5from compiler_admin.services import files 

6from compiler_admin.services.google import GoogleOrgs, GoogleUsers 

7from compiler_admin.services.toggl import GROUPS, TogglUsers 

8 

9 

10def google(format: int = Format.BASIC, inactive: bool = False, account_type: str = "", **kwargs): 

11 """Use GAM to print the users in the Google Workspace.""" 

12 if account_type and account_type not in GoogleOrgs.ORG_UNITS: 

13 raise ValueError(f"Unexpected account_type: {account_type}") 

14 

15 org_unit = GoogleOrgs.ORG_UNITS.get(account_type) 

16 org_units = [org_unit] if org_unit else [] 

17 output = GoogleUsers().get(inactive=inactive, format=format, org_units=org_units, **kwargs) 

18 click.echo(output) 

19 

20 

21def toggl(format: int = Format.BASIC, inactive: bool = False, account_type: str = "", **kwargs): 

22 """Use the Toggl API to get a list of users. 

23 

24 Mirrors the GAM output style for Google users. 

25 """ 

26 api = TogglUsers() 

27 

28 if account_type and account_type not in GROUPS: 

29 raise ValueError(f"Unexpected account_type: {account_type}") 

30 

31 group_filter = [] 

32 group_name = GROUPS.get(account_type) 

33 if group_name: 

34 group = api.get_organization_group(group_name) 

35 group_filter.append(group["group_id"]) 

36 

37 click.echo("Getting all Toggl users...", err=True) 

38 users = api.get_organization_users(inactive=inactive, groups=group_filter, **kwargs) 

39 users_df = pd.DataFrame(users) 

40 click.echo(f"Got {len(users)} Users", err=True) 

41 

42 stdout = click.get_text_stream("stdout") 

43 columns = ["email"] 

44 if format == Format.BASIC: 

45 files.write_csv(stdout, users_df, columns=columns) 

46 elif format == Format.CSV: 

47 columns += [ 

48 "name", 

49 "id", 

50 "user_id", 

51 "role_id", 

52 "organization_id", 

53 "inactive", 

54 "joined", 

55 "can_edit_email", 

56 "2fa_enabled", 

57 "avatar_url", 

58 ] 

59 files.write_csv(stdout, users_df, columns=columns) 

60 elif format == Format.JSON: 

61 files.write_json(stdout, users) 

62 

63 

64ACCOUNT_TYPES = set((*GROUPS.keys(), *GoogleOrgs.ORG_UNITS.keys())) 

65USER_SYSTEMS = {"google": google, "toggl": toggl} 

66 

67 

68@click.command() 

69@click.option( 

70 "--format", 

71 "format_key", 

72 help="The format of the output.", 

73 type=click.Choice(FORMATS.keys(), case_sensitive=False), 

74 default="basic", 

75) 

76@click.option("--inactive", help="List inactive users in the system.", is_flag=True) 

77@click.option( 

78 "-t", 

79 "--account_type", 

80 "account_type", 

81 help="Only show users with this account type.", 

82 type=click.Choice(ACCOUNT_TYPES, case_sensitive=False), 

83) 

84@click.argument( 

85 "system", 

86 type=click.Choice(USER_SYSTEMS.keys(), case_sensitive=False), 

87 default="google", 

88) 

89def users(system: str, format_key: str, inactive: bool, account_type: str, **kwargs): 

90 """List users in the Compiler workspace.""" 

91 format = FORMATS.get(format_key) 

92 

93 ls_system = USER_SYSTEMS.get(system) 

94 ls_system(format=format, inactive=inactive, account_type=account_type, **kwargs)