Skip to content

Commit

Permalink
Make WorkspaceClient and AccountClient more friendly with autospe…
Browse files Browse the repository at this point in the history
…ccing (#480)

## Changes
We cannot `create_autospec(WorkspaceClient)`, because all public
properties are created within the constructor:

```
>>> WorkspaceClient
<class 'databricks.sdk.WorkspaceClient'>
>>> ws = create_autospec(WorkspaceClient)
>>> ws
<MagicMock spec='WorkspaceClient' id='4394200816'>
>>> ws.workspace.list()
Traceback (most recent call last):
  File "/Applications/PyCharm CE.app/Contents/plugins/python-ce/helpers/pydev/_pydevd_bundle/pydevd_exec2.py", line 3, in Exec
    exec(exp, global_vars, local_vars)
  File "<input>", line 1, in <module>
  File "/opt/homebrew/Cellar/python@3.10/3.10.12_1/Frameworks/Python.framework/Versions/3.10/lib/python3.10/unittest/mock.py", line 643, in __getattr__
    raise AttributeError("Mock object has no attribute %r" % name)
AttributeError: Mock object has no attribute 'workspace'
```

This PR makes it easier to introspect the class instances and makes
mocks more reliable.

All existing behavior stays the same.

See:
- https://docs.python.org/3/library/unittest.mock.html#autospeccing
- https://stackoverflow.com/a/35921144/277035


## Tests
<!-- 
How is this tested? Please see the checklist below and also describe any
other relevant tests
-->

- [x] `make test` run locally
- [x] `make fmt` applied
- [x] relevant integration tests applied
  • Loading branch information
nfx authored Dec 20, 2023
1 parent 02bbf0c commit 1fd3b9e
Show file tree
Hide file tree
Showing 3 changed files with 615 additions and 982 deletions.
58 changes: 47 additions & 11 deletions .codegen/__init__.py.tmpl
Original file line number Diff line number Diff line change
Expand Up @@ -52,14 +52,34 @@ class WorkspaceClient:
debug_headers=debug_headers,
product=product,
product_version=product_version)
self.config = config.copy()
self.dbutils = _make_dbutils(self.config)
self.api_client = client.ApiClient(self.config)
self._config = config.copy()
self._dbutils = _make_dbutils(self._config)
self._api_client = client.ApiClient(self._config)

{{- range .Services}}{{if not .IsAccounts}}
self.{{.SnakeName}}: {{template "api" .}} = {{template "api" .}}(self.api_client){{if .Description}}
"""{{.Comment " " 100}}"""{{end}}
{{end -}}{{end}}
self._{{.SnakeName}} = {{template "api" .}}(self._api_client){{end -}}{{end}}

@property
def config(self) -> client.Config:
return self._config

@property
def api_client(self) -> client.ApiClient:
return self._api_client

@property
def dbutils(self) -> dbutils.RemoteDbUtils:
return self._dbutils

{{- range .Services}}{{if not .IsAccounts}}
@property
def {{.SnakeName}}(self) -> {{template "api" .}}:
{{if .Description}}"""{{.Summary}}"""{{end}}
return self._{{.SnakeName}}
{{end -}}{{end}}

def __repr__(self):
return f"WorkspaceClient(host='{self._config.host}', auth_type='{self._config.auth_type}', ...)"

class AccountClient:
"""
Expand All @@ -80,10 +100,26 @@ class AccountClient:
debug_headers=debug_headers,
product=product,
product_version=product_version)
self.config = config.copy()
self.api_client = client.ApiClient(self.config)
self._config = config.copy()
self._api_client = client.ApiClient(self._config)

{{- range .Services}}{{if .IsAccounts}}
self.{{(.TrimPrefix "account").SnakeName}}: {{template "api" .}} = {{template "api" .}}(self.api_client){{if .Description}}
"""{{.Comment " " 100}}"""{{end}}
{{end -}}{{end}}
self._{{(.TrimPrefix "account").SnakeName}} = {{template "api" .}}(self._api_client){{end -}}{{end}}

@property
def config(self) -> client.Config:
return self._config

@property
def api_client(self) -> client.ApiClient:
return self._api_client

{{- range .Services}}{{if .IsAccounts}}
@property
def {{(.TrimPrefix "account").SnakeName}}(self) -> {{template "api" .}}:{{if .Description}}
"""{{.Summary}}"""{{end}}
return self._{{(.TrimPrefix "account").SnakeName}}
{{end -}}{{end}}

def __repr__(self):
return f"AccountClient(account_id='{self._config.account_id}', auth_type='{self._config.auth_type}', ...)"
Loading

0 comments on commit 1fd3b9e

Please sign in to comment.