Skip to content

Commit

Permalink
ref: fix typing in sentry.testutils.factories (#67046)
Browse files Browse the repository at this point in the history
<!-- Describe your PR here. -->
  • Loading branch information
asottile-sentry authored and JonasBa committed Mar 17, 2024
1 parent 60f6309 commit ba32d5c
Show file tree
Hide file tree
Showing 5 changed files with 36 additions and 33 deletions.
1 change: 0 additions & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -485,7 +485,6 @@ module = [
"sentry.templatetags.sentry_plugins",
"sentry.testutils.asserts",
"sentry.testutils.cases",
"sentry.testutils.factories",
"sentry.testutils.fixtures",
"sentry.testutils.helpers.features",
"sentry.testutils.helpers.notifications",
Expand Down
2 changes: 1 addition & 1 deletion src/sentry/models/integrations/integration.py
Original file line number Diff line number Diff line change
Expand Up @@ -119,7 +119,7 @@ def add_organization(
"""
Add an organization to this integration.
Returns False if the OrganizationIntegration was not created
Returns None if the OrganizationIntegration was not created
"""
from sentry.models.integrations.organization_integration import OrganizationIntegration

Expand Down
62 changes: 33 additions & 29 deletions src/sentry/testutils/factories.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
import io
import os
import random
import zipfile
from base64 import b64encode
from binascii import hexlify
from collections.abc import Mapping, Sequence
Expand Down Expand Up @@ -304,22 +305,22 @@ def create_organization(name=None, owner=None, region: Region | str | None = Non
if not name:
name = petname.generate(2, " ", letters=10).title()

if region is None or SiloMode.get_current_mode() == SiloMode.MONOLITH:
region_name = get_local_region().name
org_creation_context = contextlib.nullcontext()
else:
if isinstance(region, Region):
region_name = region.name
with contextlib.ExitStack() as ctx:
if region is None or SiloMode.get_current_mode() == SiloMode.MONOLITH:
region_name = get_local_region().name
else:
region_obj = get_region_by_name(region) # Verify it exists
region_name = region_obj.name
org_creation_context = override_settings(
SILO_MODE=SiloMode.REGION, SENTRY_REGION=region_name
)
if isinstance(region, Region):
region_name = region.name
else:
region_obj = get_region_by_name(region) # Verify it exists
region_name = region_obj.name

ctx.enter_context(
override_settings(SILO_MODE=SiloMode.REGION, SENTRY_REGION=region_name)
)

with org_creation_context:
with outbox_context(flush=False):
org: Organization = Organization.objects.create(name=name, **kwargs)
org = Organization.objects.create(name=name, **kwargs)

with assume_test_silo_mode(SiloMode.CONTROL):
# Organization mapping creation relies on having a matching org slug reservation
Expand Down Expand Up @@ -592,7 +593,7 @@ def create_release(
ReleaseEnvironment.objects.create(
organization=project.organization, release=release, environment=environment
)
for project in [project] + additional_projects:
for project in [project, *additional_projects]:
ReleaseProjectEnvironment.objects.create(
project=project,
release=release,
Expand Down Expand Up @@ -626,6 +627,7 @@ def create_release(

return release

@staticmethod
def create_group_release(project: Project, group: Group, release: Release) -> GroupRelease:
return GroupRelease.objects.create(
project_id=project.id,
Expand Down Expand Up @@ -662,13 +664,11 @@ def create_release_file(release_id, file=None, name=None, dist_id=None):
def create_artifact_bundle_zip(
org=None, release=None, project=None, extra_files=None, fixture_path="artifact_bundle"
):
import zipfile

bundle = io.BytesIO()
bundle_dir = get_fixture_path(fixture_path)
with zipfile.ZipFile(bundle, "w", zipfile.ZIP_DEFLATED) as zipfile:
with zipfile.ZipFile(bundle, "w", zipfile.ZIP_DEFLATED) as zipf:
for path, content in (extra_files or {}).items():
zipfile.writestr(path, content)
zipf.writestr(path, content)
for path, _, files in os.walk(bundle_dir):
for filename in files:
fullpath = os.path.join(path, filename)
Expand All @@ -677,20 +677,20 @@ def create_artifact_bundle_zip(
manifest = _patch_artifact_manifest(
fullpath, org, release, project, extra_files
)
zipfile.writestr(relpath, manifest)
zipf.writestr(relpath, manifest)
else:
zipfile.write(fullpath, relpath)
zipf.write(fullpath, relpath)

return bundle.getvalue()

@classmethod
@assume_test_silo_mode(SiloMode.REGION)
def create_release_archive(cls, org, release: str, project=None, dist=None):
bundle = cls.create_artifact_bundle_zip(org, release, project)
file_ = File.objects.create(name="release-artifacts.zip")
file_.putfile(ContentFile(bundle))
release = Release.objects.get(organization__slug=org, version=release)
return update_artifact_index(release, dist, file_)
file = File.objects.create(name="release-artifacts.zip")
file.putfile(ContentFile(bundle))
release_obj = Release.objects.get(organization__slug=org, version=release)
return update_artifact_index(release_obj, dist, file)

@classmethod
@assume_test_silo_mode(SiloMode.REGION)
Expand Down Expand Up @@ -854,7 +854,7 @@ def create_usersocialauth(
user: User,
provider: str | None = None,
uid: str | None = None,
extra_data: Mapping[str, Any] | None = None,
extra_data: dict[str, Any] | None = None,
):
if not provider:
provider = "asana"
Expand Down Expand Up @@ -885,7 +885,7 @@ def inject_performance_problems(jobs, _):
type=group_type,
parent_span_ids=None,
cause_span_ids=None,
offender_span_ids=None,
offender_span_ids=[],
evidence_data={},
evidence_display=[],
)
Expand Down Expand Up @@ -1075,16 +1075,18 @@ def create_internal_integration_token(
) -> ApiToken:
if internal_integration and install:
raise ValueError("Only one of internal_integration or install arg can be provided")
if internal_integration is None and install is None:
elif internal_integration is None and install is None:
raise ValueError("Must pass in either internal_integration or install arg")

if install is None:
if internal_integration is not None and install is None:
# Fetch install from provided or created internal integration
with assume_test_silo_mode(SiloMode.CONTROL):
install = SentryAppInstallation.objects.get(
sentry_app=internal_integration.id,
organization_id=internal_integration.owner_id,
)
elif install is None:
raise AssertionError("unreachable")

return SentryAppInstallationTokenCreator(sentry_app_installation=install).run(
user=user, request=request
Expand Down Expand Up @@ -1138,6 +1140,7 @@ def create_sentry_app_installation(
if not prevent_token_exchange and (
install.sentry_app.status != SentryAppStatus.INTERNAL
):
assert install.api_grant is not None
GrantExchanger.run(
install=rpc_install,
code=install.api_grant.code,
Expand Down Expand Up @@ -1678,6 +1681,7 @@ def create_identity_integration(
organization_integration = integration.add_organization(
organization_id=organization.id, user=user, default_auth_id=identity.id
)
assert organization_integration is not None
return integration, organization_integration, identity, identity_provider

@staticmethod
Expand Down Expand Up @@ -1820,7 +1824,7 @@ def create_user_option(*args, **kwargs) -> UserOption:
return UserOption.objects.create(*args, **kwargs)

@staticmethod
def create_basic_auth_header(username: str, password: str = "") -> str:
def create_basic_auth_header(username: str, password: str = "") -> bytes:
return b"Basic " + b64encode(f"{username}:{password}".encode())

@staticmethod
Expand Down
2 changes: 1 addition & 1 deletion src/sentry/testutils/fixtures.py
Original file line number Diff line number Diff line change
Expand Up @@ -560,7 +560,7 @@ def create_saved_search(self, *args, **kwargs):
def create_organization_mapping(self, *args, **kwargs):
return Factories.create_org_mapping(*args, **kwargs)

def create_basic_auth_header(self, *args, **kwargs):
def create_basic_auth_header(self, *args, **kwargs) -> bytes:
return Factories.create_basic_auth_header(*args, **kwargs)

def snooze_rule(self, *args, **kwargs):
Expand Down
2 changes: 1 addition & 1 deletion src/social_auth/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,7 @@ class UserSocialAuth(models.Model):
user = models.ForeignKey(AUTH_USER_MODEL, related_name="social_auth", on_delete=models.CASCADE)
provider = models.CharField(max_length=32)
uid = models.CharField(max_length=UID_LENGTH)
extra_data: models.Field[dict[str, Any], dict[str, Any]] = JSONField(default="{}")
extra_data: models.Field[dict[str, Any] | None, dict[str, Any]] = JSONField(default="{}")

class Meta:
"""Meta data"""
Expand Down

0 comments on commit ba32d5c

Please sign in to comment.