<commit_before># Copyright (c) 2016 Ultimaker B.V.
# Uranium is released under the terms of the AGPLv3 or higher.

import pytest
import os.path

from UM.MimeTypeDatabase import MimeType, MimeTypeDatabase

@pytest.fixture
def mime_database():
    MimeTypeDatabase._MimeTypeDatabase__custom_mimetypes = []
    mime = MimeType(
        name = "application/x-test",
        comment = "Test Mimetype",
        suffixes = [ "test" ]
    )
    MimeTypeDatabase.addMimeType(mime)
    mime = MimeType(
        name = "application/x-long-test",
        comment = "Long Test Mimetype",
        suffixes = [ "long.test" ]
    )
    MimeTypeDatabase.addMimeType(mime)
    return MimeTypeDatabase

def test_system_mimetypes(mime_database):
    mime = mime_database.getMimeTypeForFile(os.path.abspath(__file__))
    assert mime.name == "text/x-python"
    assert mime.comment == "Python script"
    assert mime.suffixes == ["py", "pyx", "wsgi"]
    assert mime.preferredSuffix == "py"

def test_compare(mime_database):
    mime1 = mime_database.getMimeType("text/plain")
    mime2 = mime_database.getMimeType("text/plain")

    assert id(mime1) != id(mime2)
    assert mime1 == mime2

##  Tests the creation of new MIME types.
def test_createMimeType():
    # Normal MIME type
    mime = MimeType(
        name = "application/x-test",
        comment = "Normal MIME type",
        suffixes = [ "tst" ]
    )
    assert mime.name == "application/x-test"
    assert mime.comment == "Normal MIME type"
    assert "tst" in mime.suffixes

    # No name.
    with pytest.raises(ValueError):
        mime = MimeType(
            name = None,
            comment = "Missing name",
            suffixes = [ "non" ]
        )

    # No suffixes.
    mime = MimeType(
        name = "application/x-no-suffixes",
        comment = "No Suffixes",
        suffixes = []
    )
    assert len(mime.suffixes) == 0

    # Multiple suffixes.
    mime = MimeType(
        name = "application/x-multiple",
        comment = "Multiple suffixes",
        suffixes = [ "boo", "baa" ]
    )
    assert "boo" in mime.suffixes
    assert "baa" in mime.suffixes

    # A preferred suffix.
    mime = MimeType(
        name = "application/x-preferred",
        comment = "Preferred suffix",
        suffixes = [ "boo", "baa" ],
        preferred_suffix = "baa"
    )
    assert mime.preferredSuffix == "baa"

    # Nonexistent preference.
    with pytest.raises(ValueError):
        mime = MimeType(
            name = "application/x-nonexistent-pref",
            comment = "Nonexistent preferred suffix",
            suffixes = [ "boo", "baa" ],
            preferred_suffix = "bee" # Not in the list of suffixes.
        )

def test_custom_mimetypes(mime_database):
    mime = mime_database.getMimeTypeForFile(os.path.join(os.path.dirname(os.path.abspath(__file__)), "file.test"))
    assert mime.name == "application/x-test"

    mime = mime_database.getMimeTypeForFile(os.path.join(os.path.dirname(os.path.abspath(__file__)), "file.long.test"))
    assert mime.name == "application/x-long-test"

def test_stripExtension(mime_database):
    mime = mime_database.getMimeTypeForFile(os.path.join(os.path.dirname(os.path.abspath(__file__)), "file.test"))

    assert mime.stripExtension("file.test") == "file"
    assert mime.stripExtension("file.long.test") == "file.long"
    assert mime.stripExtension("some.random.file.txt") == "some.random.file.txt"

    mime = mime_database.getMimeTypeForFile(os.path.join(os.path.dirname(os.path.abspath(__file__)), "file.long.test"))

    assert mime.stripExtension("file.test") == "file.test"
    assert mime.stripExtension("file.long.test") == "file"
    assert mime.stripExtension("some.random.file.txt") == "some.random.file.txt"

<commit_msg>Expand MIME type test with multiple extension case<commit_after># Copyright (c) 2016 Ultimaker B.V.
# Uranium is released under the terms of the AGPLv3 or higher.

import pytest
import os.path

from UM.MimeTypeDatabase import MimeType, MimeTypeDatabase

@pytest.fixture
def mime_database():
    MimeTypeDatabase._MimeTypeDatabase__custom_mimetypes = []
    mime = MimeType(
        name = "application/x-test",
        comment = "Test Mimetype",
        suffixes = [ "test" ]
    )
    MimeTypeDatabase.addMimeType(mime)
    mime = MimeType(
        name = "application/x-long-test",
        comment = "Long Test Mimetype",
        suffixes = [ "long.test" ]
    )
    MimeTypeDatabase.addMimeType(mime)
    mime = MimeType(
        name = "application/x-multiple-test",
        comment = "Multiple Extension MIME type",
        suffixes = [ "tost", "ost" ],
        preferred_suffix = "ost"
    )
    MimeTypeDatabase.addMimeType(mime)
    return MimeTypeDatabase

def test_system_mimetypes(mime_database):
    mime = mime_database.getMimeTypeForFile(os.path.abspath(__file__))
    assert mime.name == "text/x-python"
    assert mime.comment == "Python script"
    assert mime.suffixes == ["py", "pyx", "wsgi"]
    assert mime.preferredSuffix == "py"

def test_compare(mime_database):
    mime1 = mime_database.getMimeType("text/plain")
    mime2 = mime_database.getMimeType("text/plain")

    assert id(mime1) != id(mime2)
    assert mime1 == mime2

##  Tests the creation of new MIME types.
def test_createMimeType():
    # Normal MIME type
    mime = MimeType(
        name = "application/x-test",
        comment = "Normal MIME type",
        suffixes = [ "tst" ]
    )
    assert mime.name == "application/x-test"
    assert mime.comment == "Normal MIME type"
    assert "tst" in mime.suffixes

    # No name.
    with pytest.raises(ValueError):
        mime = MimeType(
            name = None,
            comment = "Missing name",
            suffixes = [ "non" ]
        )

    # No suffixes.
    mime = MimeType(
        name = "application/x-no-suffixes",
        comment = "No Suffixes",
        suffixes = []
    )
    assert len(mime.suffixes) == 0

    # Multiple suffixes.
    mime = MimeType(
        name = "application/x-multiple",
        comment = "Multiple suffixes",
        suffixes = [ "boo", "baa" ]
    )
    assert "boo" in mime.suffixes
    assert "baa" in mime.suffixes

    # A preferred suffix.
    mime = MimeType(
        name = "application/x-preferred",
        comment = "Preferred suffix",
        suffixes = [ "boo", "baa" ],
        preferred_suffix = "baa"
    )
    assert mime.preferredSuffix == "baa"

    # Nonexistent preference.
    with pytest.raises(ValueError):
        mime = MimeType(
            name = "application/x-nonexistent-pref",
            comment = "Nonexistent preferred suffix",
            suffixes = [ "boo", "baa" ],
            preferred_suffix = "bee" # Not in the list of suffixes.
        )

def test_custom_mimetypes(mime_database):
    mime = mime_database.getMimeTypeForFile(os.path.join(os.path.dirname(os.path.abspath(__file__)), "file.test"))
    assert mime.name == "application/x-test"

    mime = mime_database.getMimeTypeForFile(os.path.join(os.path.dirname(os.path.abspath(__file__)), "file.long.test"))
    assert mime.name == "application/x-long-test"

def test_stripExtension(mime_database):
    mime = mime_database.getMimeTypeForFile(os.path.join(os.path.dirname(os.path.abspath(__file__)), "file.test"))
    assert mime.stripExtension("file.test") == "file"
    assert mime.stripExtension("file.long.test") == "file.long"
    assert mime.stripExtension("some.random.file.txt") == "some.random.file.txt"

    mime = mime_database.getMimeTypeForFile(os.path.join(os.path.dirname(os.path.abspath(__file__)), "file.long.test"))
    assert mime.stripExtension("file.test") == "file.test"
    assert mime.stripExtension("file.long.test") == "file"
    assert mime.stripExtension("some.random.file.txt") == "some.random.file.txt"

    mime = mime_database.getMimeTypeForFile(os.path.join(os.path.dirname(os.path.abspath(__file__)), "file.tost"))
    assert mime.stripExtension("file.tost") == "file"
    assert mime.stripExtension("file.ost") == "file"
    assert mime.stripExtension("file.tost.png") == "file.tost.png"<|endoftext|>


<commit_before>"""
BSD 3-Clause License

Copyright (c) 2017, Mairie de Paris
All rights reserved.

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:

* Redistributions of source code must retain the above copyright notice, this
  list of conditions and the following disclaimer.

* Redistributions in binary form must reproduce the above copyright notice,
  this list of conditions and the following disclaimer in the documentation
  and/or other materials provided with the distribution.

* Neither the name of the copyright holder nor the names of its
  contributors may be used to endorse or promote products derived from
  this software without specific prior written permission.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
"""

import os

DEBUG = True
HOST = os.getenv('HOST', '0.0.0.0')
PORT = int(os.getenv('PORT', '5000'))

# Grabs the folder where the script runs.
BASEDIR = os.path.abspath(os.path.dirname(__file__))

UPLOAD_FOLDER = os.path.abspath(BASEDIR + "/../uploads")
ALLOWED_MIME = set([
    'image/jpeg',
    'image/png',
    'application/pdf',
])
MAX_CONTENT_LENGTH = 5 * 1024 * 1024

KEEP_SCANS = False
<commit_msg>KEEP_SCANS use env var<commit_after>"""
BSD 3-Clause License

Copyright (c) 2017, Mairie de Paris
All rights reserved.

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:

* Redistributions of source code must retain the above copyright notice, this
  list of conditions and the following disclaimer.

* Redistributions in binary form must reproduce the above copyright notice,
  this list of conditions and the following disclaimer in the documentation
  and/or other materials provided with the distribution.

* Neither the name of the copyright holder nor the names of its
  contributors may be used to endorse or promote products derived from
  this software without specific prior written permission.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
"""

import os

DEBUG = os.getenv("DEBUG", False)
HOST = os.getenv("HOST", "0.0.0.0")
PORT = int(os.getenv("PORT", "5000"))

# Grabs the folder where the script runs.
BASEDIR = os.path.abspath(os.path.dirname(__file__))

UPLOAD_FOLDER = os.path.abspath(BASEDIR + "/../uploads")
ALLOWED_MIME = set([
    "image/jpeg",
    "image/png",
    "application/pdf",
])
MAX_CONTENT_LENGTH = 5 * 1024 * 1024

KEEP_SCANS = os.getenv("KEEP_SCANS", False)
<|endoftext|>


<commit_before>import os

if os.environ.get('VCAP_APPLICATION'):
    # on cloudfoundry, config is a json blob in VCAP_APPLICATION - unpack it, and populate
    # standard environment variables from it
    from app.cloudfoundry_config import extract_cloudfoundry_config
    extract_cloudfoundry_config()

class Config(object):
    ADMIN_CLIENT_SECRET = os.environ.get('ADMIN_CLIENT_SECRET')
    API_HOST_NAME = os.environ.get('API_HOST_NAME')
    SECRET_KEY = os.environ.get('SECRET_KEY')
    DANGEROUS_SALT = os.environ.get('DANGEROUS_SALT')
    ZENDESK_API_KEY = os.environ.get('ZENDESK_API_KEY')

    # if we're not on cloudfoundry, we can get to this app from localhost. but on cloudfoundry its different
    ADMIN_BASE_URL = os.environ.get('ADMIN_BASE_URL', 'http://localhost:6012')

    TEMPLATE_PREVIEW_API_HOST = os.environ.get('TEMPLATE_PREVIEW_API_HOST', 'http://localhost:6013')
    TEMPLATE_PREVIEW_API_KEY = os.environ.get('TEMPLATE_PREVIEW_API_KEY', 'my-secret-key')

    # Logging
    DEBUG = False
    NOTIFY_LOG_PATH = os.getenv('NOTIFY_LOG_PATH')

    ADMIN_CLIENT_USER_NAME = 'notify-admin'

    ANTIVIRUS_API_HOST = os.environ.get('ANTIVIRUS_API_HOST')
    ANTIVIRUS_API_KEY = os.environ.get('ANTIVIRUS_API_KEY')

    ASSETS_DEBUG = False
    AWS_REGION = 'eu-west-1'
    DEFAULT_SERVICE_LIMIT = 50
    DEFAULT_FREE_SMS_FRAGMENT_LIMITS = {
        'central': 250000,
        'local': 25000,
        'nhs_central': 250000,
        'nhs_local': 25000,
        'nhs_gp': 25000,
        'emergency_service': 25000,
        'school_or_college': 25000,
        'other': 25000,
    }
    EMAIL_EXPIRY_SECONDS = 3600  # 1 hour
    INVITATION_EXPIRY_SECONDS = 3600 * 24 * 2  # 2 days - also set on api
    EMAIL_2FA_EXPIRY_SECONDS = 1800  # 30 Minutes
    HEADER_COLOUR = '#81878b'  # mix(govuk-colour("dark-grey"), govuk-colour("mid-grey"))
    HTTP_PROTOCOL = 'http'
    MAX_FAILED_LOGIN_COUNT = 10
    NOTIFY_APP_NAME = 'admin'
    NOTIFY_LOG_LEVEL = 'DEBUG'
    PERMANENT_SESSION_LIFETIME = 20 * 60 * 60  # 20 hours
    SEND_FILE_MAX_AGE_DEFAULT = 365 * 24 * 60 * 60  # 1 year
    SESSION_COOKIE_HTTPONLY = True
    SESSION_COOKIE_NAME = 'notify_admin_session'
    SESSION_COOKIE_SECURE = True
    # don't send back the cookie if it hasn't been modified by the request. this means that the expiry time won't be
    # updated unless the session is changed - but it's generally refreshed by `save_service_or_org_after_request`
    # every time anyway, except for specific endpoints (png/pdfs generally) where we've disabled that handler.
    SESSION_REFRESH_EACH_REQUEST = False
    SHOW_STYLEGUIDE = True
    WTF_CSRF_ENABLED = True
    WTF_CSRF_TIME_LIMIT = None
    CSV_UPLOAD_BUCKET_NAME = 'local-notifications-csv-upload'
    CONTACT_LIST_UPLOAD_BUCKET_NAME = 'local-contact-list'
    ACTIVITY_STATS_LIMIT_DAYS = 7
    TEST_MESSAGE_FILENAME = 'Report'

    REPLY_TO_EMAIL_ADDRESS_VALIDATION_TIMEOUT = 45

    NOTIFY_ENVIRONMENT = 'development'
    LOGO_UPLOAD_BUCKET_NAME = 'public-logos-local'
    MOU_BUCKET_NAME = 'local-mou'
    TRANSIENT_UPLOADED_LETTERS = 'local-transient-uploaded-letters'
    ROUTE_SECRET_KEY_1 = os.environ.get('ROUTE_SECRET_KEY_1', '')
    ROUTE_SECRET_KEY_2 = os.environ.get('ROUTE_SECRET_KEY_2', '')
    CHECK_PROXY_HEADER = False
    ANTIVIRUS_ENABLED = True

    REDIS_URL = os.environ.get('REDIS_URL')
    REDIS_ENABLED = os.environ.get('REDIS_ENABLED') == '1'

    ASSET_DOMAIN = ''
    ASSET_PATH = '/static/'

    # as defined in api db migration 0331_add_broadcast_org.py
    BROADCAST_ORGANISATION_ID = '38e4bf69-93b0-445d-acee-53ea53fe02df'

    NOTIFY_SERVICE_ID = 'd6aa2c68-a2d9-4437-ab19-3ae8eb202553'

class Development(Config):
    NOTIFY_LOG_PATH = 'application.log'
    DEBUG = True
    SESSION_COOKIE_SECURE = False
    SESSION_PROTECTION = None
    CSV_UPLOAD_BUCKET_NAME = 'development-notifications-csv-upload'
    CONTACT_LIST_UPLOAD_BUCKET_NAME = 'development-contact-list'
    LOGO_UPLOAD_BUCKET_NAME = 'public-logos-tools'
    MOU_BUCKET_NAME = 'notify.tools-mou'
    TRANSIENT_UPLOADED_LETTERS = 'development-transient-uploaded-letters'

    ADMIN_CLIENT_SECRET = 'dev-notify-secret-key'
    API_HOST_NAME = 'http://localhost:6011'
    DANGEROUS_SALT = 'dev-notify-salt'
    SECRET_KEY = 'dev-notify-secret-key'
    ANTIVIRUS_API_HOST = 'http://localhost:6016'
    ANTIVIRUS_API_KEY = 'test-key'
    ANTIVIRUS_ENABLED = os.getenv('ANTIVIRUS_ENABLED') == '1'

    ASSET_PATH = '/static/'

    REDIS_URL = 'redis://localhost:6379/0'

class Test(Development):
    DEBUG = True
    TESTING = True
    WTF_CSRF_ENABLED = False
    CSV_UPLOAD_BUCKET_NAME = 'test-notifications-csv-upload'
    CONTACT_LIST_UPLOAD_BUCKET_NAME = 'test-contact-list'
    LOGO_UPLOAD_BUCKET_NAME = 'public-logos-test'
    MOU_BUCKET_NAME = 'test-mou'
    TRANSIENT_UPLOADED_LETTERS = 'test-transient-uploaded-letters'
    NOTIFY_ENVIRONMENT = 'test'
    API_HOST_NAME = 'http://you-forgot-to-mock-an-api-call-to'
    TEMPLATE_PREVIEW_API_HOST = 'http://localhost:9999'
    ANTIVIRUS_API_HOST = 'https://test-antivirus'
    ANTIVIRUS_API_KEY = 'test-antivirus-secret'
    ANTIVIRUS_ENABLED = True

    ASSET_DOMAIN = 'static.example.com'
    ASSET_PATH = 'https://static.example.com/'

class Preview(Config):
    HTTP_PROTOCOL = 'https'
    HEADER_COLOUR = '#F499BE'  # $baby-pink
    CSV_UPLOAD_BUCKET_NAME = 'preview-notifications-csv-upload'
    CONTACT_LIST_UPLOAD_BUCKET_NAME = 'preview-contact-list'
    LOGO_UPLOAD_BUCKET_NAME = 'public-logos-preview'
    MOU_BUCKET_NAME = 'notify.works-mou'
    TRANSIENT_UPLOADED_LETTERS = 'preview-transient-uploaded-letters'
    NOTIFY_ENVIRONMENT = 'preview'
    CHECK_PROXY_HEADER = False
    ASSET_DOMAIN = 'static.notify.works'
    ASSET_PATH = 'https://static.notify.works/'

    # On preview, extend the validation timeout to allow more leniency when running functional tests
    REPLY_TO_EMAIL_ADDRESS_VALIDATION_TIMEOUT = 120

class Staging(Config):
    SHOW_STYLEGUIDE = False
    HTTP_PROTOCOL = 'https'
    HEADER_COLOUR = '#6F72AF'  # $mauve
    CSV_UPLOAD_BUCKET_NAME = 'staging-notifications-csv-upload'
    CONTACT_LIST_UPLOAD_BUCKET_NAME = 'staging-contact-list'
    LOGO_UPLOAD_BUCKET_NAME = 'public-logos-staging'
    MOU_BUCKET_NAME = 'staging-notify.works-mou'
    TRANSIENT_UPLOADED_LETTERS = 'staging-transient-uploaded-letters'
    NOTIFY_ENVIRONMENT = 'staging'
    CHECK_PROXY_HEADER = False
    ASSET_DOMAIN = 'static.staging-notify.works'
    ASSET_PATH = 'https://static.staging-notify.works/'

class Live(Config):
    SHOW_STYLEGUIDE = False
    HEADER_COLOUR = '#005EA5'  # $govuk-blue
    HTTP_PROTOCOL = 'https'
    CSV_UPLOAD_BUCKET_NAME = 'live-notifications-csv-upload'
    CONTACT_LIST_UPLOAD_BUCKET_NAME = 'production-contact-list'
    LOGO_UPLOAD_BUCKET_NAME = 'public-logos-production'
    MOU_BUCKET_NAME = 'notifications.service.gov.uk-mou'
    TRANSIENT_UPLOADED_LETTERS = 'production-transient-uploaded-letters'
    NOTIFY_ENVIRONMENT = 'live'
    CHECK_PROXY_HEADER = False
    ASSET_DOMAIN = 'static.notifications.service.gov.uk'
    ASSET_PATH = 'https://static.notifications.service.gov.uk/'

class CloudFoundryConfig(Config):
    pass

# CloudFoundry sandbox
class Sandbox(CloudFoundryConfig):
    HTTP_PROTOCOL = 'https'
    HEADER_COLOUR = '#F499BE'  # $baby-pink
    CSV_UPLOAD_BUCKET_NAME = 'cf-sandbox-notifications-csv-upload'
    LOGO_UPLOAD_BUCKET_NAME = 'cf-sandbox-notifications-logo-upload'
    NOTIFY_ENVIRONMENT = 'sandbox'

configs = {
    'development': Development,
    'test': Test,
    'preview': Preview,
    'staging': Staging,
    'live': Live,
    'production': Live,
    'sandbox': Sandbox
}
<commit_msg>Remove reference to styleguide from config<commit_after>import os

if os.environ.get('VCAP_APPLICATION'):
    # on cloudfoundry, config is a json blob in VCAP_APPLICATION - unpack it, and populate
    # standard environment variables from it
    from app.cloudfoundry_config import extract_cloudfoundry_config
    extract_cloudfoundry_config()

class Config(object):
    ADMIN_CLIENT_SECRET = os.environ.get('ADMIN_CLIENT_SECRET')
    API_HOST_NAME = os.environ.get('API_HOST_NAME')
    SECRET_KEY = os.environ.get('SECRET_KEY')
    DANGEROUS_SALT = os.environ.get('DANGEROUS_SALT')
    ZENDESK_API_KEY = os.environ.get('ZENDESK_API_KEY')

    # if we're not on cloudfoundry, we can get to this app from localhost. but on cloudfoundry its different
    ADMIN_BASE_URL = os.environ.get('ADMIN_BASE_URL', 'http://localhost:6012')

    TEMPLATE_PREVIEW_API_HOST = os.environ.get('TEMPLATE_PREVIEW_API_HOST', 'http://localhost:6013')
    TEMPLATE_PREVIEW_API_KEY = os.environ.get('TEMPLATE_PREVIEW_API_KEY', 'my-secret-key')

    # Logging
    DEBUG = False
    NOTIFY_LOG_PATH = os.getenv('NOTIFY_LOG_PATH')

    ADMIN_CLIENT_USER_NAME = 'notify-admin'

    ANTIVIRUS_API_HOST = os.environ.get('ANTIVIRUS_API_HOST')
    ANTIVIRUS_API_KEY = os.environ.get('ANTIVIRUS_API_KEY')

    ASSETS_DEBUG = False
    AWS_REGION = 'eu-west-1'
    DEFAULT_SERVICE_LIMIT = 50
    DEFAULT_FREE_SMS_FRAGMENT_LIMITS = {
        'central': 250000,
        'local': 25000,
        'nhs_central': 250000,
        'nhs_local': 25000,
        'nhs_gp': 25000,
        'emergency_service': 25000,
        'school_or_college': 25000,
        'other': 25000,
    }
    EMAIL_EXPIRY_SECONDS = 3600  # 1 hour
    INVITATION_EXPIRY_SECONDS = 3600 * 24 * 2  # 2 days - also set on api
    EMAIL_2FA_EXPIRY_SECONDS = 1800  # 30 Minutes
    HEADER_COLOUR = '#81878b'  # mix(govuk-colour("dark-grey"), govuk-colour("mid-grey"))
    HTTP_PROTOCOL = 'http'
    MAX_FAILED_LOGIN_COUNT = 10
    NOTIFY_APP_NAME = 'admin'
    NOTIFY_LOG_LEVEL = 'DEBUG'
    PERMANENT_SESSION_LIFETIME = 20 * 60 * 60  # 20 hours
    SEND_FILE_MAX_AGE_DEFAULT = 365 * 24 * 60 * 60  # 1 year
    SESSION_COOKIE_HTTPONLY = True
    SESSION_COOKIE_NAME = 'notify_admin_session'
    SESSION_COOKIE_SECURE = True
    # don't send back the cookie if it hasn't been modified by the request. this means that the expiry time won't be
    # updated unless the session is changed - but it's generally refreshed by `save_service_or_org_after_request`
    # every time anyway, except for specific endpoints (png/pdfs generally) where we've disabled that handler.
    SESSION_REFRESH_EACH_REQUEST = False
    WTF_CSRF_ENABLED = True
    WTF_CSRF_TIME_LIMIT = None
    CSV_UPLOAD_BUCKET_NAME = 'local-notifications-csv-upload'
    CONTACT_LIST_UPLOAD_BUCKET_NAME = 'local-contact-list'
    ACTIVITY_STATS_LIMIT_DAYS = 7
    TEST_MESSAGE_FILENAME = 'Report'

    REPLY_TO_EMAIL_ADDRESS_VALIDATION_TIMEOUT = 45

    NOTIFY_ENVIRONMENT = 'development'
    LOGO_UPLOAD_BUCKET_NAME = 'public-logos-local'
    MOU_BUCKET_NAME = 'local-mou'
    TRANSIENT_UPLOADED_LETTERS = 'local-transient-uploaded-letters'
    ROUTE_SECRET_KEY_1 = os.environ.get('ROUTE_SECRET_KEY_1', '')
    ROUTE_SECRET_KEY_2 = os.environ.get('ROUTE_SECRET_KEY_2', '')
    CHECK_PROXY_HEADER = False
    ANTIVIRUS_ENABLED = True

    REDIS_URL = os.environ.get('REDIS_URL')
    REDIS_ENABLED = os.environ.get('REDIS_ENABLED') == '1'

    ASSET_DOMAIN = ''
    ASSET_PATH = '/static/'

    # as defined in api db migration 0331_add_broadcast_org.py
    BROADCAST_ORGANISATION_ID = '38e4bf69-93b0-445d-acee-53ea53fe02df'

    NOTIFY_SERVICE_ID = 'd6aa2c68-a2d9-4437-ab19-3ae8eb202553'

class Development(Config):
    NOTIFY_LOG_PATH = 'application.log'
    DEBUG = True
    SESSION_COOKIE_SECURE = False
    SESSION_PROTECTION = None
    CSV_UPLOAD_BUCKET_NAME = 'development-notifications-csv-upload'
    CONTACT_LIST_UPLOAD_BUCKET_NAME = 'development-contact-list'
    LOGO_UPLOAD_BUCKET_NAME = 'public-logos-tools'
    MOU_BUCKET_NAME = 'notify.tools-mou'
    TRANSIENT_UPLOADED_LETTERS = 'development-transient-uploaded-letters'

    ADMIN_CLIENT_SECRET = 'dev-notify-secret-key'
    API_HOST_NAME = 'http://localhost:6011'
    DANGEROUS_SALT = 'dev-notify-salt'
    SECRET_KEY = 'dev-notify-secret-key'
    ANTIVIRUS_API_HOST = 'http://localhost:6016'
    ANTIVIRUS_API_KEY = 'test-key'
    ANTIVIRUS_ENABLED = os.getenv('ANTIVIRUS_ENABLED') == '1'

    ASSET_PATH = '/static/'

    REDIS_URL = 'redis://localhost:6379/0'

class Test(Development):
    DEBUG = True
    TESTING = True
    WTF_CSRF_ENABLED = False
    CSV_UPLOAD_BUCKET_NAME = 'test-notifications-csv-upload'
    CONTACT_LIST_UPLOAD_BUCKET_NAME = 'test-contact-list'
    LOGO_UPLOAD_BUCKET_NAME = 'public-logos-test'
    MOU_BUCKET_NAME = 'test-mou'
    TRANSIENT_UPLOADED_LETTERS = 'test-transient-uploaded-letters'
    NOTIFY_ENVIRONMENT = 'test'
    API_HOST_NAME = 'http://you-forgot-to-mock-an-api-call-to'
    TEMPLATE_PREVIEW_API_HOST = 'http://localhost:9999'
    ANTIVIRUS_API_HOST = 'https://test-antivirus'
    ANTIVIRUS_API_KEY = 'test-antivirus-secret'
    ANTIVIRUS_ENABLED = True

    ASSET_DOMAIN = 'static.example.com'
    ASSET_PATH = 'https://static.example.com/'

class Preview(Config):
    HTTP_PROTOCOL = 'https'
    HEADER_COLOUR = '#F499BE'  # $baby-pink
    CSV_UPLOAD_BUCKET_NAME = 'preview-notifications-csv-upload'
    CONTACT_LIST_UPLOAD_BUCKET_NAME = 'preview-contact-list'
    LOGO_UPLOAD_BUCKET_NAME = 'public-logos-preview'
    MOU_BUCKET_NAME = 'notify.works-mou'
    TRANSIENT_UPLOADED_LETTERS = 'preview-transient-uploaded-letters'
    NOTIFY_ENVIRONMENT = 'preview'
    CHECK_PROXY_HEADER = False
    ASSET_DOMAIN = 'static.notify.works'
    ASSET_PATH = 'https://static.notify.works/'

    # On preview, extend the validation timeout to allow more leniency when running functional tests
    REPLY_TO_EMAIL_ADDRESS_VALIDATION_TIMEOUT = 120

class Staging(Config):
    HTTP_PROTOCOL = 'https'
    HEADER_COLOUR = '#6F72AF'  # $mauve
    CSV_UPLOAD_BUCKET_NAME = 'staging-notifications-csv-upload'
    CONTACT_LIST_UPLOAD_BUCKET_NAME = 'staging-contact-list'
    LOGO_UPLOAD_BUCKET_NAME = 'public-logos-staging'
    MOU_BUCKET_NAME = 'staging-notify.works-mou'
    TRANSIENT_UPLOADED_LETTERS = 'staging-transient-uploaded-letters'
    NOTIFY_ENVIRONMENT = 'staging'
    CHECK_PROXY_HEADER = False
    ASSET_DOMAIN = 'static.staging-notify.works'
    ASSET_PATH = 'https://static.staging-notify.works/'

class Live(Config):
    HEADER_COLOUR = '#005EA5'  # $govuk-blue
    HTTP_PROTOCOL = 'https'
    CSV_UPLOAD_BUCKET_NAME = 'live-notifications-csv-upload'
    CONTACT_LIST_UPLOAD_BUCKET_NAME = 'production-contact-list'
    LOGO_UPLOAD_BUCKET_NAME = 'public-logos-production'
    MOU_BUCKET_NAME = 'notifications.service.gov.uk-mou'
    TRANSIENT_UPLOADED_LETTERS = 'production-transient-uploaded-letters'
    NOTIFY_ENVIRONMENT = 'live'
    CHECK_PROXY_HEADER = False
    ASSET_DOMAIN = 'static.notifications.service.gov.uk'
    ASSET_PATH = 'https://static.notifications.service.gov.uk/'

class CloudFoundryConfig(Config):
    pass

# CloudFoundry sandbox
class Sandbox(CloudFoundryConfig):
    HTTP_PROTOCOL = 'https'
    HEADER_COLOUR = '#F499BE'  # $baby-pink
    CSV_UPLOAD_BUCKET_NAME = 'cf-sandbox-notifications-csv-upload'
    LOGO_UPLOAD_BUCKET_NAME = 'cf-sandbox-notifications-logo-upload'
    NOTIFY_ENVIRONMENT = 'sandbox'

configs = {
    'development': Development,
    'test': Test,
    'preview': Preview,
    'staging': Staging,
    'live': Live,
    'production': Live,
    'sandbox': Sandbox
}
<|endoftext|>


<commit_before>"""
I/O for the STL format, cf.
<https://en.wikipedia.org/wiki/STL_(file_format)>.
"""
import logging
import os

import numpy

from ..__about__ import __version__
from .._exceptions import ReadError, WriteError
from .._files import open_file
from .._helpers import register
from .._mesh import CellBlock, Mesh

def read(filename):
    with open_file(filename, "rb") as f:
        # Checking if the file is ASCII format is normally done by checking if the
        # first 5 characters of the header is "solid".
        # ```
        # header = f.read(80).decode("utf-8")
        # ```
        # Unfortunately, there are mesh files out there which are binary and still put
        # "solid" there.
        # A suggested alternative is to do as if the file is binary, read the
        # num_triangles and see if it matches the file size
        # (https://stackoverflow.com/a/7394842/353337).
        f.read(80)
        num_triangles = numpy.fromfile(f, count=1, dtype=numpy.uint32)[0]
        # for each triangle, one has 3 float32 (facet normal), 9 float32 (facet), and 1
        # int16 (attribute count), 50 bytes in total
        is_binary = 84 + num_triangles * 50 == os.path.getsize(filename)
        if is_binary:
            out = _read_binary(f, num_triangles)
        else:
            # skip header
            f.seek(0)
            f.readline()
            out = _read_ascii(f)
    return out

# numpy.loadtxt is super slow
# Code adapted from <https://stackoverflow.com/a/8964779/353337>.
def iter_loadtxt(infile, skiprows=0, comments=["#"], dtype=float, usecols=None):
    def iter_func():
        for _ in range(skiprows):
            next(infile)
        for line in infile:
            line = line.decode("utf-8").strip()
            if line.startswith(comments):
                continue
            # remove all text
            items = line.split()[-3:]
            usecols_ = range(len(items)) if usecols is None else usecols
            for idx in usecols_:
                yield dtype(items[idx])
        iter_loadtxt.rowlength = len(items) if usecols is None else len(usecols)

    data = numpy.fromiter(iter_func(), dtype=dtype)
    return data.reshape((-1, iter_loadtxt.rowlength))

def _read_ascii(f):
    # The file has the form
    # ```
    # solid foo
    #   facet normal 0.455194 -0.187301 -0.870469
    #    outer loop
    #     vertex 266.36 234.594 14.6145
    #     vertex 268.582 234.968 15.6956
    #     vertex 267.689 232.646 15.7283
    #    endloop
    #   endfacet
    #   # [...] more facets [...]
    # endsolid
    # ```
    # In the interest of speed, don't verify the format and instead just skip the text.

    # TODO Pandas is MUCH faster than numpy for i/o, see
    # <https://stackoverflow.com/a/18260092/353337>.
    # import pandas
    # data = pandas.read_csv(
    #     f,
    #     skiprows=lambda row: row == 0 or (row - 1) % 7 in [0, 1, 5, 6],
    #     skipfooter=1,
    #     usecols=(1, 2, 3),
    # )

    # numpy.loadtxt is super slow
    # data = numpy.loadtxt(
    #     f,
    #     comments=["solid", "facet", "outer loop", "endloop", "endfacet", "endsolid"],
    #     usecols=(1, 2, 3),
    # )
    data = iter_loadtxt(
        f,
        comments=("solid", "outer loop", "endloop", "endfacet", "endsolid"),
        # usecols=(1, 2, 3),
    )

    if data.shape[0] % 4 != 0:
        raise ReadError()

    # split off the facet normals
    facet_rows = numpy.zeros(len(data), dtype=bool)
    facet_rows[0::4] = True
    facet_normals = data[facet_rows]
    data = data[~facet_rows]

    facets = numpy.split(data, data.shape[0] // 3)
    points, cells = data_from_facets(facets)
    return Mesh(points, cells, cell_data={"facet_normals": [facet_normals]})

def data_from_facets(facets):
    # Now, all facets contain the point coordinate. Try to identify individual
    # points and build the data arrays.
    pts = numpy.concatenate(facets)

    # TODO equip `unique()` with a tolerance
    # Use return_index so we can use sort on `idx` such that the order is
    # preserved; see <https://stackoverflow.com/a/15637512/353337>.
    _, idx, inv = numpy.unique(pts, axis=0, return_index=True, return_inverse=True)
    k = numpy.argsort(idx)
    points = pts[idx[k]]
    inv_k = numpy.argsort(k)
    cells = [CellBlock("triangle", inv_k[inv].reshape(-1, 3))]
    return points, cells

def _read_binary(f, num_triangles):
    # for each triangle, one has 3 float32 (facet normal), 9 float32 (facet), and 1
    # int16 (attribute count)
    out = numpy.fromfile(
        f,
        count=num_triangles,
        dtype=numpy.dtype(
            [("normal", "f4", (3,)), ("facet", "f4", (3, 3)), ("attr count", "i2")]
        ),
    )
    # discard normals, attribute count
    facets = out["facet"]
    # if not numpy.all(out["attr count"] == 0):
    #     print(out["attr count"])
    #     raise ReadError("Nonzero attr count")

    points, cells = data_from_facets(facets)
    return Mesh(points, cells)

def write(filename, mesh, binary=False):
    if not any(c.type == "triangle" for c in mesh.cells):
        raise WriteError(
            "STL can only write triangle cells (not {}).".format(
                ", ".join(c.type for c in mesh.cells)
            )
        )
    if len(mesh.cells) > 1:
        keys = set(mesh.cells.keys())
        keys.remove("triangle")
        logging.warning(
            "STL can only write triangle cells. Discarding {}.".format(", ".join(keys))
        )

    if mesh.points.shape[1] == 2:
        logging.warning(
            "STL requires 3D points, but 2D points given. "
            "Appending 0 third component."
        )
        mesh.points = numpy.column_stack(
            [mesh.points[:, 0], mesh.points[:, 1], numpy.zeros(mesh.points.shape[0])]
        )

    if binary:
        _binary(filename, mesh.points, mesh.cells)
    else:
        _write_ascii(filename, mesh)

def _compute_normals(pts):
    normals = numpy.cross(pts[:, 1] - pts[:, 0], pts[:, 2] - pts[:, 0])
    nrm = numpy.sqrt(numpy.einsum("ij,ij->i", normals, normals))
    normals = (normals.T / nrm).T
    return normals

def _write_ascii(filename, mesh):
    with open_file(filename, "wb") as fh:
        fh.write(b"solid\n")
        pts = mesh.points[mesh.get_cells_type("triangle")]
        if "facet_normals" in mesh.cell_data:
            normals = mesh.get_cell_data("facet_normals", "triangle")
        else:
            normals = _compute_normals(pts)

        for local_pts, normal in zip(pts, normals):
            # facet normal 0.455194 -0.187301 -0.870469
            #  outer loop
            #   vertex 266.36 234.594 14.6145
            #   vertex 268.582 234.968 15.6956
            #   vertex 267.689 232.646 15.7283
            #  endloop
            # endfacet
            out = ["facet normal {} {} {}".format(*normal), " outer loop"]
            for pt in local_pts:
                out += ["  vertex {} {} {}".format(*pt)]
            out += [" endloop", "endfacet"]
            fh.write(("\n".join(out) + "\n").encode("utf-8"))
        fh.write(b"endsolid\n")

def _binary(filename, points, cells):
    with open_file(filename, "wb") as fh:
        # 80 character header data
        msg = "This file was generated by meshio v{}.".format(__version__)
        msg += (79 - len(msg)) * "X"
        msg += "\n"
        fh.write(msg.encode("utf-8"))
        for c in filter(lambda c: c.type == "triangle", cells):
            pts = points[c.data]
            normals = _compute_normals(pts)
            fh.write(numpy.uint32(len(c.data)))
            for pt, normal in zip(pts, normals):
                fh.write(normal.astype(numpy.float32))
                fh.write(pt.astype(numpy.float32))
                fh.write(numpy.uint16(0))

register("stl", [".stl"], read, {"stl": write})
<commit_msg>fix stl error<commit_after>"""
I/O for the STL format, cf.
<https://en.wikipedia.org/wiki/STL_(file_format)>.
"""
import logging
import os

import numpy

from ..__about__ import __version__
from .._exceptions import ReadError, WriteError
from .._files import open_file
from .._helpers import register
from .._mesh import CellBlock, Mesh

def read(filename):
    with open_file(filename, "rb") as f:
        # Checking if the file is ASCII format is normally done by checking if the
        # first 5 characters of the header is "solid".
        # ```
        # header = f.read(80).decode("utf-8")
        # ```
        # Unfortunately, there are mesh files out there which are binary and still put
        # "solid" there.
        # A suggested alternative is to do as if the file is binary, read the
        # num_triangles and see if it matches the file size
        # (https://stackoverflow.com/a/7394842/353337).
        f.read(80)
        num_triangles = numpy.fromfile(f, count=1, dtype=numpy.uint32)[0]
        # for each triangle, one has 3 float32 (facet normal), 9 float32 (facet), and 1
        # int16 (attribute count), 50 bytes in total
        is_binary = 84 + num_triangles * 50 == os.path.getsize(filename)
        if is_binary:
            out = _read_binary(f, num_triangles)
        else:
            # skip header
            f.seek(0)
            f.readline()
            out = _read_ascii(f)
    return out

# numpy.loadtxt is super slow
# Code adapted from <https://stackoverflow.com/a/8964779/353337>.
def iter_loadtxt(infile, skiprows=0, comments=["#"], dtype=float, usecols=None):
    def iter_func():
        for _ in range(skiprows):
            next(infile)
        for line in infile:
            line = line.decode("utf-8").strip()
            if line.startswith(comments):
                continue
            # remove all text
            items = line.split()[-3:]
            usecols_ = range(len(items)) if usecols is None else usecols
            for idx in usecols_:
                yield dtype(items[idx])
        iter_loadtxt.rowlength = len(items) if usecols is None else len(usecols)

    data = numpy.fromiter(iter_func(), dtype=dtype)
    return data.reshape((-1, iter_loadtxt.rowlength))

def _read_ascii(f):
    # The file has the form
    # ```
    # solid foo
    #   facet normal 0.455194 -0.187301 -0.870469
    #    outer loop
    #     vertex 266.36 234.594 14.6145
    #     vertex 268.582 234.968 15.6956
    #     vertex 267.689 232.646 15.7283
    #    endloop
    #   endfacet
    #   # [...] more facets [...]
    # endsolid
    # ```
    # In the interest of speed, don't verify the format and instead just skip the text.

    # TODO Pandas is MUCH faster than numpy for i/o, see
    # <https://stackoverflow.com/a/18260092/353337>.
    # import pandas
    # data = pandas.read_csv(
    #     f,
    #     skiprows=lambda row: row == 0 or (row - 1) % 7 in [0, 1, 5, 6],
    #     skipfooter=1,
    #     usecols=(1, 2, 3),
    # )

    # numpy.loadtxt is super slow
    # data = numpy.loadtxt(
    #     f,
    #     comments=["solid", "facet", "outer loop", "endloop", "endfacet", "endsolid"],
    #     usecols=(1, 2, 3),
    # )
    data = iter_loadtxt(
        f,
        comments=("solid", "outer loop", "endloop", "endfacet", "endsolid"),
        # usecols=(1, 2, 3),
    )

    if data.shape[0] % 4 != 0:
        raise ReadError()

    # split off the facet normals
    facet_rows = numpy.zeros(len(data), dtype=bool)
    facet_rows[0::4] = True
    facet_normals = data[facet_rows]
    data = data[~facet_rows]

    facets = numpy.split(data, data.shape[0] // 3)
    points, cells = data_from_facets(facets)
    return Mesh(points, cells, cell_data={"facet_normals": [facet_normals]})

def data_from_facets(facets):
    # Now, all facets contain the point coordinate. Try to identify individual
    # points and build the data arrays.
    pts = numpy.concatenate(facets)

    # TODO equip `unique()` with a tolerance
    # Use return_index so we can use sort on `idx` such that the order is
    # preserved; see <https://stackoverflow.com/a/15637512/353337>.
    _, idx, inv = numpy.unique(pts, axis=0, return_index=True, return_inverse=True)
    k = numpy.argsort(idx)
    points = pts[idx[k]]
    inv_k = numpy.argsort(k)
    cells = [CellBlock("triangle", inv_k[inv].reshape(-1, 3))]
    return points, cells

def _read_binary(f, num_triangles):
    # for each triangle, one has 3 float32 (facet normal), 9 float32 (facet), and 1
    # int16 (attribute count)
    out = numpy.fromfile(
        f,
        count=num_triangles,
        dtype=numpy.dtype(
            [("normal", "f4", (3,)), ("facet", "f4", (3, 3)), ("attr count", "i2")]
        ),
    )
    # discard normals, attribute count
    facets = out["facet"]
    # if not numpy.all(out["attr count"] == 0):
    #     print(out["attr count"])
    #     raise ReadError("Nonzero attr count")

    points, cells = data_from_facets(facets)
    return Mesh(points, cells)

def write(filename, mesh, binary=False):
    if not any(c.type == "triangle" for c in mesh.cells):
        raise WriteError(
            "STL can only write triangle cells (not {}).".format(
                ", ".join(c.type for c in mesh.cells)
            )
        )
    if len(mesh.cells) > 1:
        invalid = set(block.type for block in mesh.cells if block.type != "triangle")
        logging.warning(
            "STL can only write triangle cells. Discarding {}.".format(", ".join(invalid))
        )

    if mesh.points.shape[1] == 2:
        logging.warning(
            "STL requires 3D points, but 2D points given. "
            "Appending 0 third component."
        )
        mesh.points = numpy.column_stack(
            [mesh.points[:, 0], mesh.points[:, 1], numpy.zeros(mesh.points.shape[0])]
        )

    if binary:
        _binary(filename, mesh.points, mesh.cells)
    else:
        _write_ascii(filename, mesh)

def _compute_normals(pts):
    normals = numpy.cross(pts[:, 1] - pts[:, 0], pts[:, 2] - pts[:, 0])
    nrm = numpy.sqrt(numpy.einsum("ij,ij->i", normals, normals))
    normals = (normals.T / nrm).T
    return normals

def _write_ascii(filename, mesh):
    with open_file(filename, "wb") as fh:
        fh.write(b"solid\n")
        pts = mesh.points[mesh.get_cells_type("triangle")]
        if "facet_normals" in mesh.cell_data:
            normals = mesh.get_cell_data("facet_normals", "triangle")
        else:
            normals = _compute_normals(pts)

        for local_pts, normal in zip(pts, normals):
            # facet normal 0.455194 -0.187301 -0.870469
            #  outer loop
            #   vertex 266.36 234.594 14.6145
            #   vertex 268.582 234.968 15.6956
            #   vertex 267.689 232.646 15.7283
            #  endloop
            # endfacet
            out = ["facet normal {} {} {}".format(*normal), " outer loop"]
            for pt in local_pts:
                out += ["  vertex {} {} {}".format(*pt)]
            out += [" endloop", "endfacet"]
            fh.write(("\n".join(out) + "\n").encode("utf-8"))
        fh.write(b"endsolid\n")

def _binary(filename, points, cells):
    with open_file(filename, "wb") as fh:
        # 80 character header data
        msg = "This file was generated by meshio v{}.".format(__version__)
        msg += (79 - len(msg)) * "X"
        msg += "\n"
        fh.write(msg.encode("utf-8"))
        for c in filter(lambda c: c.type == "triangle", cells):
            pts = points[c.data]
            normals = _compute_normals(pts)
            fh.write(numpy.uint32(len(c.data)))
            for pt, normal in zip(pts, normals):
                fh.write(normal.astype(numpy.float32))
                fh.write(pt.astype(numpy.float32))
                fh.write(numpy.uint16(0))

register("stl", [".stl"], read, {"stl": write})
<|endoftext|>


<commit_before># Copyright 2013 OpenStack Foundation
# All rights reserved.
#
#    Licensed under the Apache License, Version 2.0 (the "License"); you may
#    not use this file except in compliance with the License. You may obtain
#    a copy of the License at
#
#         http://www.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
#    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
#    License for the specific language governing permissions and limitations
#    under the License.

"""
ML2 Mechanism Driver for Cisco Nexus platforms.
"""

from oslo.config import cfg

from neutron.common import constants as n_const
from neutron.extensions import portbindings
from neutron.openstack.common import log as logging
from neutron.plugins.common import constants as p_const
from neutron.plugins.ml2 import driver_api as api
from neutron.plugins.ml2.drivers.cisco.nexus import config as conf
from neutron.plugins.ml2.drivers.cisco.nexus import credentials_v2 as cred
from neutron.plugins.ml2.drivers.cisco.nexus import exceptions as excep
from neutron.plugins.ml2.drivers.cisco.nexus import nexus_db_v2 as nxos_db
from neutron.plugins.ml2.drivers.cisco.nexus import nexus_network_driver

LOG = logging.getLogger(__name__)

class CiscoNexusMechanismDriver(api.MechanismDriver):

    """Cisco Nexus ML2 Mechanism Driver."""

    def initialize(self):
        # Create ML2 device dictionary from ml2_conf.ini entries.
        conf.ML2MechCiscoConfig()

        # Extract configuration parameters from the configuration file.
        self._nexus_switches = conf.ML2MechCiscoConfig.nexus_dict
        LOG.debug(_("nexus_switches found = %s"), self._nexus_switches)

        self.credentials = {}
        self.driver = nexus_network_driver.CiscoNexusDriver()

        # Initialize credential store after database initialization.
        cred.Store.initialize()

    def _valid_network_segment(self, segment):
        return (cfg.CONF.ml2_cisco.managed_physical_network is None or
                cfg.CONF.ml2_cisco.managed_physical_network ==
                segment[api.PHYSICAL_NETWORK])

    def _get_vlanid(self, segment):
        if (segment and segment[api.NETWORK_TYPE] == p_const.TYPE_VLAN and
            self._valid_network_segment(segment)):
            return segment.get(api.SEGMENTATION_ID)

    def _is_deviceowner_compute(self, port):
        return port['device_owner'].startswith('compute')

    def _is_status_active(self, port):
        return port['status'] == n_const.PORT_STATUS_ACTIVE

    def _get_switch_info(self, host_id):
        for switch_ip, attr in self._nexus_switches:
            if str(attr) == str(host_id):
                port_id = self._nexus_switches[switch_ip, attr]
                return port_id, switch_ip
        else:
            raise excep.NexusComputeHostNotConfigured(host=host_id)

    def _configure_nxos_db(self, vlan_id, device_id, host_id):
        """Create the nexus database entry.

        Called during update precommit port event.
        """
        port_id, switch_ip = self._get_switch_info(host_id)
        nxos_db.add_nexusport_binding(port_id, str(vlan_id), switch_ip,
                                      device_id)

    def _configure_switch_entry(self, vlan_id, device_id, host_id):
        """Create a nexus switch entry.

        if needed, create a VLAN in the appropriate switch/port and
        configure the appropriate interfaces for this VLAN.

        Called during update postcommit port event.
        """
        port_id, switch_ip = self._get_switch_info(host_id)
        vlan_name = cfg.CONF.ml2_cisco.vlan_name_prefix + str(vlan_id)

        # Check to see if this is the first binding to use this vlan on the
        # switch/port. Configure switch accordingly.
        bindings = nxos_db.get_nexusvlan_binding(vlan_id, switch_ip)
        if len(bindings) == 1:
            LOG.debug(_("Nexus: create & trunk vlan %s"), vlan_name)
            self.driver.create_and_trunk_vlan(switch_ip, vlan_id, vlan_name,
                                              port_id)
        else:
            LOG.debug(_("Nexus: trunk vlan %s"), vlan_name)
            self.driver.enable_vlan_on_trunk_int(switch_ip, vlan_id, port_id)

    def _delete_nxos_db(self, vlan_id, device_id, host_id):
        """Delete the nexus database entry.

        Called during delete precommit port event.
        """
        try:
            row = nxos_db.get_nexusvm_binding(vlan_id, device_id)
            nxos_db.remove_nexusport_binding(row.port_id, row.vlan_id,
                                             row.switch_ip, row.instance_id)
        except excep.NexusPortBindingNotFound:
            return

    def _delete_switch_entry(self, vlan_id, device_id, host_id):
        """Delete the nexus switch entry.

        By accessing the current db entries determine if switch
        configuration can be removed.

        Called during update postcommit port event.
        """
        port_id, switch_ip = self._get_switch_info(host_id)

        # if there are no remaining db entries using this vlan on this nexus
        # switch port then remove vlan from the switchport trunk.
        try:
            nxos_db.get_port_vlan_switch_binding(port_id, vlan_id, switch_ip)
        except excep.NexusPortBindingNotFound:
            self.driver.disable_vlan_on_trunk_int(switch_ip, vlan_id, port_id)

            # if there are no remaining db entries using this vlan on this
            # nexus switch then remove the vlan.
            try:
                nxos_db.get_nexusvlan_binding(vlan_id, switch_ip)
            except excep.NexusPortBindingNotFound:
                self.driver.delete_vlan(switch_ip, vlan_id)

    def _is_vm_migration(self, context):
        if not context.bound_segment and context.original_bound_segment:
            return (context.current.get(portbindings.HOST_ID) !=
                    context.original.get(portbindings.HOST_ID))

    def _port_action(self, port, segment, func):
        """Verify configuration and then process event."""
        device_id = port.get('device_id')
        host_id = port.get(portbindings.HOST_ID)

        # Workaround until vlan can be retrieved during delete_port_postcommit
        # event.
        if func == self._delete_switch_entry:
            vlan_id = self._delete_port_postcommit_vlan
        else:
            vlan_id = self._get_vlanid(segment)

        if vlan_id and device_id and host_id:
            func(vlan_id, device_id, host_id)
        else:
            fields = "vlan_id " if not vlan_id else ""
            fields += "device_id " if not device_id else ""
            fields += "host_id" if not host_id else ""
            raise excep.NexusMissingRequiredFields(fields=fields)

        # Workaround until vlan can be retrieved during delete_port_postcommit
        # event.
        if func == self._delete_nxos_db:
            self._delete_port_postcommit_vlan = vlan_id
        else:
            self._delete_port_postcommit_vlan = 0

    def update_port_precommit(self, context):
        """Update port pre-database transaction commit event."""

        # if VM migration is occurring then remove previous database entry
        # else process update event.
        if self._is_vm_migration(context):
            self._port_action(context.original,
                              context.original_bound_segment,
                              self._delete_nxos_db)
        else:
            if (self._is_deviceowner_compute(context.current) and
                self._is_status_active(context.current)):
                self._port_action(context.current,
                                  context.bound_segment,
                                  self._configure_nxos_db)

    def update_port_postcommit(self, context):
        """Update port non-database commit event."""

        # if VM migration is occurring then remove previous nexus switch entry
        # else process update event.
        if self._is_vm_migration(context):
            self._port_action(context.original,
                              context.original_bound_segment,
                              self._delete_switch_entry)
        else:
            if (self._is_deviceowner_compute(context.current) and
                self._is_status_active(context.current)):
                self._port_action(context.current,
                                  context.bound_segment,
                                  self._configure_switch_entry)

    def delete_port_precommit(self, context):
        """Delete port pre-database commit event."""
        if self._is_deviceowner_compute(context.current):
            self._port_action(context.current,
                              context.bound_segment,
                              self._delete_nxos_db)

    def delete_port_postcommit(self, context):
        """Delete port non-database commit event."""
        if self._is_deviceowner_compute(context.current):
            self._port_action(context.current,
                              context.bound_segment,
                              self._delete_switch_entry)
<commit_msg>ML2 Cisco Nexus MD: Remove workaround for bug 1276395<commit_after># Copyright 2013 OpenStack Foundation
# All rights reserved.
#
#    Licensed under the Apache License, Version 2.0 (the "License"); you may
#    not use this file except in compliance with the License. You may obtain
#    a copy of the License at
#
#         http://www.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
#    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
#    License for the specific language governing permissions and limitations
#    under the License.

"""
ML2 Mechanism Driver for Cisco Nexus platforms.
"""

from oslo.config import cfg

from neutron.common import constants as n_const
from neutron.extensions import portbindings
from neutron.openstack.common import log as logging
from neutron.plugins.common import constants as p_const
from neutron.plugins.ml2 import driver_api as api
from neutron.plugins.ml2.drivers.cisco.nexus import config as conf
from neutron.plugins.ml2.drivers.cisco.nexus import credentials_v2 as cred
from neutron.plugins.ml2.drivers.cisco.nexus import exceptions as excep
from neutron.plugins.ml2.drivers.cisco.nexus import nexus_db_v2 as nxos_db
from neutron.plugins.ml2.drivers.cisco.nexus import nexus_network_driver

LOG = logging.getLogger(__name__)

class CiscoNexusMechanismDriver(api.MechanismDriver):

    """Cisco Nexus ML2 Mechanism Driver."""

    def initialize(self):
        # Create ML2 device dictionary from ml2_conf.ini entries.
        conf.ML2MechCiscoConfig()

        # Extract configuration parameters from the configuration file.
        self._nexus_switches = conf.ML2MechCiscoConfig.nexus_dict
        LOG.debug(_("nexus_switches found = %s"), self._nexus_switches)

        self.credentials = {}
        self.driver = nexus_network_driver.CiscoNexusDriver()

        # Initialize credential store after database initialization.
        cred.Store.initialize()

    def _valid_network_segment(self, segment):
        return (cfg.CONF.ml2_cisco.managed_physical_network is None or
                cfg.CONF.ml2_cisco.managed_physical_network ==
                segment[api.PHYSICAL_NETWORK])

    def _get_vlanid(self, segment):
        if (segment and segment[api.NETWORK_TYPE] == p_const.TYPE_VLAN and
            self._valid_network_segment(segment)):
            return segment.get(api.SEGMENTATION_ID)

    def _is_deviceowner_compute(self, port):
        return port['device_owner'].startswith('compute')

    def _is_status_active(self, port):
        return port['status'] == n_const.PORT_STATUS_ACTIVE

    def _get_switch_info(self, host_id):
        for switch_ip, attr in self._nexus_switches:
            if str(attr) == str(host_id):
                port_id = self._nexus_switches[switch_ip, attr]
                return port_id, switch_ip
        else:
            raise excep.NexusComputeHostNotConfigured(host=host_id)

    def _configure_nxos_db(self, vlan_id, device_id, host_id):
        """Create the nexus database entry.

        Called during update precommit port event.
        """
        port_id, switch_ip = self._get_switch_info(host_id)
        nxos_db.add_nexusport_binding(port_id, str(vlan_id), switch_ip,
                                      device_id)

    def _configure_switch_entry(self, vlan_id, device_id, host_id):
        """Create a nexus switch entry.

        if needed, create a VLAN in the appropriate switch/port and
        configure the appropriate interfaces for this VLAN.

        Called during update postcommit port event.
        """
        port_id, switch_ip = self._get_switch_info(host_id)
        vlan_name = cfg.CONF.ml2_cisco.vlan_name_prefix + str(vlan_id)

        # Check to see if this is the first binding to use this vlan on the
        # switch/port. Configure switch accordingly.
        bindings = nxos_db.get_nexusvlan_binding(vlan_id, switch_ip)
        if len(bindings) == 1:
            LOG.debug(_("Nexus: create & trunk vlan %s"), vlan_name)
            self.driver.create_and_trunk_vlan(switch_ip, vlan_id, vlan_name,
                                              port_id)
        else:
            LOG.debug(_("Nexus: trunk vlan %s"), vlan_name)
            self.driver.enable_vlan_on_trunk_int(switch_ip, vlan_id, port_id)

    def _delete_nxos_db(self, vlan_id, device_id, host_id):
        """Delete the nexus database entry.

        Called during delete precommit port event.
        """
        try:
            row = nxos_db.get_nexusvm_binding(vlan_id, device_id)
            nxos_db.remove_nexusport_binding(row.port_id, row.vlan_id,
                                             row.switch_ip, row.instance_id)
        except excep.NexusPortBindingNotFound:
            return

    def _delete_switch_entry(self, vlan_id, device_id, host_id):
        """Delete the nexus switch entry.

        By accessing the current db entries determine if switch
        configuration can be removed.

        Called during update postcommit port event.
        """
        port_id, switch_ip = self._get_switch_info(host_id)

        # if there are no remaining db entries using this vlan on this nexus
        # switch port then remove vlan from the switchport trunk.
        try:
            nxos_db.get_port_vlan_switch_binding(port_id, vlan_id, switch_ip)
        except excep.NexusPortBindingNotFound:
            self.driver.disable_vlan_on_trunk_int(switch_ip, vlan_id, port_id)

            # if there are no remaining db entries using this vlan on this
            # nexus switch then remove the vlan.
            try:
                nxos_db.get_nexusvlan_binding(vlan_id, switch_ip)
            except excep.NexusPortBindingNotFound:
                self.driver.delete_vlan(switch_ip, vlan_id)

    def _is_vm_migration(self, context):
        if not context.bound_segment and context.original_bound_segment:
            return (context.current.get(portbindings.HOST_ID) !=
                    context.original.get(portbindings.HOST_ID))

    def _port_action(self, port, segment, func):
        """Verify configuration and then process event."""
        device_id = port.get('device_id')
        host_id = port.get(portbindings.HOST_ID)
        vlan_id = self._get_vlanid(segment)

        if vlan_id and device_id and host_id:
            func(vlan_id, device_id, host_id)
        else:
            fields = "vlan_id " if not vlan_id else ""
            fields += "device_id " if not device_id else ""
            fields += "host_id" if not host_id else ""
            raise excep.NexusMissingRequiredFields(fields=fields)

    def update_port_precommit(self, context):
        """Update port pre-database transaction commit event."""

        # if VM migration is occurring then remove previous database entry
        # else process update event.
        if self._is_vm_migration(context):
            self._port_action(context.original,
                              context.original_bound_segment,
                              self._delete_nxos_db)
        else:
            if (self._is_deviceowner_compute(context.current) and
                self._is_status_active(context.current)):
                self._port_action(context.current,
                                  context.bound_segment,
                                  self._configure_nxos_db)

    def update_port_postcommit(self, context):
        """Update port non-database commit event."""

        # if VM migration is occurring then remove previous nexus switch entry
        # else process update event.
        if self._is_vm_migration(context):
            self._port_action(context.original,
                              context.original_bound_segment,
                              self._delete_switch_entry)
        else:
            if (self._is_deviceowner_compute(context.current) and
                self._is_status_active(context.current)):
                self._port_action(context.current,
                                  context.bound_segment,
                                  self._configure_switch_entry)

    def delete_port_precommit(self, context):
        """Delete port pre-database commit event."""
        if self._is_deviceowner_compute(context.current):
            self._port_action(context.current,
                              context.bound_segment,
                              self._delete_nxos_db)

    def delete_port_postcommit(self, context):
        """Delete port non-database commit event."""
        if self._is_deviceowner_compute(context.current):
            self._port_action(context.current,
                              context.bound_segment,
                              self._delete_switch_entry)
<|endoftext|>


<commit_before>#!/usr/bin/env python

import six

from csvkit.exceptions import ColumnIdentifierError

class FilteringCSVReader(six.Iterator):
    """
    Given any row iterator, only return rows which pass the filter.
    If 'header' is False, then all rows must pass the filter; by default, the first row will be passed
    through untested.

    The value of patterns may be either a sequence or a dictionary.  Items in the sequence and values in the
    dictionary may be strings, regular expressions, or functions.  For each row in the wrapped iterator,
    these values will be used as tests, and the row will only be yielded by the filter if all values pass
    their corresponding tests.  This behavior can be toggled so that all rows which pass any of the tests
    will be yielded by specifying "any_match=True" in the constructor.

    Empty values (the blank string or None) not be tested; the value in that position will not affect whether
    or not the filtering reader yields a prospective row.  To test for explicitly blank, use a regular
    expression such as "^$" or "^\s*$"

    If patterns is a dictionary, the keys can be integers identifying indices in the input rows, or, if 'header'
    is True (as it is by default), they can be strings matching column names in the first row of the reader.

    If patterns is a sequence, then it is assumed that they will be applied to the
    equivalently positioned values in the test rows.

    By specifying 'inverse=True', only rows which do not match the patterns will be passed by the filter. The header,
    if there is one, will always be returned regardless of the value for 'inverse'.
    """
    returned_header = False
    column_names = None

    def __init__(self, reader, patterns, header=True, any_match=False, inverse=False):
        super(FilteringCSVReader, self).__init__()

        self.reader = reader
        self.header = header

        if self.header:
            self.column_names = next(reader)

        self.any_match = any_match
        self.inverse = inverse
        self.patterns = standardize_patterns(self.column_names, patterns)

    def __iter__(self):
        return self

    def __next__(self):
        if self.column_names and not self.returned_header:
            self.returned_header = True
            return self.column_names

        while True:
            row = next(self.reader)

            if self.test_row(row):
                return row

        raise StopIteration()

    def test_row(self, row):
        for idx, test in self.patterns.items():
            try:
                value = row[idx]
            except IndexError:
                value = ''
            result = test(value)
            if self.any_match:
                if result:
                    return not self.inverse  # True
            else:
                if not result:
                    return self.inverse  # False

        if self.any_match:
            return self.inverse  # False
        else:
            return not self.inverse  # True

def standardize_patterns(column_names, patterns):
    """
    Given patterns in any of the permitted input forms, return a dict whose keys
    are column indices and whose values are functions which return a boolean value whether the value passes.
    If patterns is a dictionary and any of its keys are values in column_names, the returned dictionary will
    have those keys replaced with the integer position of that value in column_names
    """
    try:
        # Dictionary of patterns
        patterns = dict((k, pattern_as_function(v)) for k, v in patterns.items() if v)
        if not column_names:
            return patterns
        p2 = {}
        for k in patterns:
            if k in column_names:
                idx = column_names.index(k)
                if idx in patterns:
                    raise ColumnIdentifierError("Column %s has index %i which already has a pattern." % (k, idx))
                p2[idx] = patterns[k]
            else:
                p2[k] = patterns[k]
        return p2
    except AttributeError:
        # Sequence of patterns
        return dict((i, pattern_as_function(x)) for i, x in enumerate(patterns))

def pattern_as_function(obj):
    # obj is function
    if hasattr(obj, '__call__'):
        return obj

    # obj is regex object
    if hasattr(obj, 'match'):
        return regex_callable(obj)

    # obj is string
    return lambda x: obj in x

class regex_callable(object):

    def __init__(self, pattern):
        self.pattern = pattern

    def __call__(self, arg):
        return self.pattern.search(arg)
<commit_msg>Fix warning because of unescaped '\s'<commit_after>#!/usr/bin/env python

import six

from csvkit.exceptions import ColumnIdentifierError

class FilteringCSVReader(six.Iterator):
    r"""
    Given any row iterator, only return rows which pass the filter.
    If 'header' is False, then all rows must pass the filter; by default, the first row will be passed
    through untested.

    The value of patterns may be either a sequence or a dictionary.  Items in the sequence and values in the
    dictionary may be strings, regular expressions, or functions.  For each row in the wrapped iterator,
    these values will be used as tests, and the row will only be yielded by the filter if all values pass
    their corresponding tests.  This behavior can be toggled so that all rows which pass any of the tests
    will be yielded by specifying "any_match=True" in the constructor.

    Empty values (the blank string or None) not be tested; the value in that position will not affect whether
    or not the filtering reader yields a prospective row.  To test for explicitly blank, use a regular
    expression such as "^$" or "^\s*$"

    If patterns is a dictionary, the keys can be integers identifying indices in the input rows, or, if 'header'
    is True (as it is by default), they can be strings matching column names in the first row of the reader.

    If patterns is a sequence, then it is assumed that they will be applied to the
    equivalently positioned values in the test rows.

    By specifying 'inverse=True', only rows which do not match the patterns will be passed by the filter. The header,
    if there is one, will always be returned regardless of the value for 'inverse'.
    """
    returned_header = False
    column_names = None

    def __init__(self, reader, patterns, header=True, any_match=False, inverse=False):
        super(FilteringCSVReader, self).__init__()

        self.reader = reader
        self.header = header

        if self.header:
            self.column_names = next(reader)

        self.any_match = any_match
        self.inverse = inverse
        self.patterns = standardize_patterns(self.column_names, patterns)

    def __iter__(self):
        return self

    def __next__(self):
        if self.column_names and not self.returned_header:
            self.returned_header = True
            return self.column_names

        while True:
            row = next(self.reader)

            if self.test_row(row):
                return row

        raise StopIteration()

    def test_row(self, row):
        for idx, test in self.patterns.items():
            try:
                value = row[idx]
            except IndexError:
                value = ''
            result = test(value)
            if self.any_match:
                if result:
                    return not self.inverse  # True
            else:
                if not result:
                    return self.inverse  # False

        if self.any_match:
            return self.inverse  # False
        else:
            return not self.inverse  # True

def standardize_patterns(column_names, patterns):
    """
    Given patterns in any of the permitted input forms, return a dict whose keys
    are column indices and whose values are functions which return a boolean value whether the value passes.
    If patterns is a dictionary and any of its keys are values in column_names, the returned dictionary will
    have those keys replaced with the integer position of that value in column_names
    """
    try:
        # Dictionary of patterns
        patterns = dict((k, pattern_as_function(v)) for k, v in patterns.items() if v)
        if not column_names:
            return patterns
        p2 = {}
        for k in patterns:
            if k in column_names:
                idx = column_names.index(k)
                if idx in patterns:
                    raise ColumnIdentifierError("Column %s has index %i which already has a pattern." % (k, idx))
                p2[idx] = patterns[k]
            else:
                p2[k] = patterns[k]
        return p2
    except AttributeError:
        # Sequence of patterns
        return dict((i, pattern_as_function(x)) for i, x in enumerate(patterns))

def pattern_as_function(obj):
    # obj is function
    if hasattr(obj, '__call__'):
        return obj

    # obj is regex object
    if hasattr(obj, 'match'):
        return regex_callable(obj)

    # obj is string
    return lambda x: obj in x

class regex_callable(object):

    def __init__(self, pattern):
        self.pattern = pattern

    def __call__(self, arg):
        return self.pattern.search(arg)
<|endoftext|>


<commit_before># Copyright 2014 Google Inc. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Runs fio benchmarks to simulate logging, database and streaming.

Man: http://manpages.ubuntu.com/manpages/natty/man1/fio.1.html
Quick howto: http://www.bluestop.org/fio/HOWTO.txt

Simulated logging benchmark does the following things (PD only):
0) Do NOT use direct IO for any tests below, simply go through the FS.
1) Sequentially write x GB with queue depth equal to 8, where x is decided by
   the test VM's total memory. (A larger VM will write more bytes)
2) Random read of 10% of the bytes written.
3) Sequential read of all of the bytes written.

Simulated database benchmark does the following things (PD, PD-SSD, local SSD):
1) 4K Random R on a file using queue depths 1, 16 and 64 (each queue depth
   is a different benchmark).
2) 4K Random W on a file using queue depths 1, 16 and 64 (each queue depth
   is a different benchmark).
3) 4K Random 90% R/ 10% W on a file using queue depths 1, 16 and 64 (each
   queue depth is a different benchmark).
4) The size of the test file is decided by the test VM's total memory and capped
   at 1GB to ensure this test finishes within reasonable time.

Simulated streaming benchmark (PD only):
1) 1M Seq R at queue depth 1 and 16 (streaming).
2) 1M Seq W at queue depth 1 and 16 (streaming).

For AWS, where use PD, we should use EBS-GP and EBS Magnetic, for PD-SSD use
EBS-GP and PIOPS.
"""

import logging
import re

from perfkitbenchmarker import flags

LOGGING = 'logging'
DATABASE = 'database'
STREAMING = 'streaming'

flags.DEFINE_enum('workload_mode', LOGGING,
                  [LOGGING, DATABASE, STREAMING],
                  'Simulate a logging, database or streaming scenario.')

flags.DEFINE_list('iodepth_list', None, 'A list of iodepth parameter used by '
                  'fio command in simulated database and streaming scenarios '
                  'only.')

flags.DEFINE_integer('maxjobs', 0,
                     'The maximum allowed number of jobs to support.')

FLAGS = flags.FLAGS

BENCHMARK_INFO = {'name': 'synthetic_storage_workloads_benchmark',
                  'description': 'Runs FIO in sequenctial, random, read and '
                                 'write modes to simulate various scenarios.',
                  'scratch_disk': True,
                  'num_machines': 1}
REQUIRED_PACKAGES = 'bc fio libaio1'
DESCRIPTION = 'description'
METHOD = 'method'

DEFAULT_IODEPTH = 8
DEFAULT_DATABASE_SIMULATION_IODEPTH_LIST = [1, 16, 64]
DEFAULT_STREAMING_SIMULATION_IODEPTH_LIST = [1, 16]

LATENCY_REGEX = r'[=\s]+([\d\.]+)[\s,]+'
BANDWIDTH_REGEX = r'(\d+)(\w+/*\w*)'

def GetInfo():
  return BENCHMARK_INFO

def Prepare(benchmark_spec):
  """Prepare the virtual machine to run FIO.

     This includes installing fio, bc. and libaio1 and insuring that the
     attached disk is large enough to support the fio benchmark.

  Args:
    benchmark_spec: The benchmark specification. Contains all data that is
        required to run the benchmark.
  """
  vms = benchmark_spec.vms
  vm = vms[0]
  logging.info('FIO prepare on %s', vm)
  vm.InstallPackage(REQUIRED_PACKAGES)

def ParseFioResult(res):
  """Parse result from fio commands.

  Args:
    res: Result from fio commands in string format.
  Returns:
    A list of tuples represents latency. Each tuple contains unit and
    latency number.
    A list of tuples represents bandwidth. Each tuple contains aggrb,
    unit of aggrb, minb, unit of minb, maxb, unit of maxb, mint, unit of mint,
    maxt, unit of maxt.
  """
  latency = re.findall(r'\s+lat \((\w+)\)[\s:]+'
                       r'min%smax%savg%sstdev%s' % (
                           LATENCY_REGEX, LATENCY_REGEX,
                           LATENCY_REGEX, LATENCY_REGEX), res)
  bandwidth = re.findall(
      r'aggrb=%s, minb=%s, maxb=%s, mint=%s, maxt=%s' % (
          BANDWIDTH_REGEX, BANDWIDTH_REGEX, BANDWIDTH_REGEX, BANDWIDTH_REGEX,
          BANDWIDTH_REGEX), res)
  print latency, bandwidth
  return latency, bandwidth

def CreateSampleFromBandwidthTuple(result, test, iodepth=DEFAULT_IODEPTH):
  """Create a sample from bandwidth result tuple.

  Args:
    result: A tuple, containing (aggrb, unit of aggrb, minb, unit of minb,
        maxb, unit of maxb, mint, unit of mint, maxt, unit of maxt).
    test: Name of test.
    iodepth: Iodepth parameter used by fio command.
  Returns:
    A list of samples in the form of 3 or 4 tuples. The tuples contain
        the sample metric (string), value (float), and unit (string).
        If a 4th element is included, it is a dictionary of sample
        metadata.
  """
  return [test, float(result[0]), result[1],
          {'minb': result[2] + result[3],
           'maxb': result[4] + result[5],
           'mint': result[6] + result[7],
           'maxt': result[8] + result[9],
           'max_jobs': FLAGS.maxjobs,
           'iodepth': iodepth}]

def CreateSampleFromLatencyTuple(result, test, iodepth=DEFAULT_IODEPTH):
  """Create a sample from latency result tuple.

  Args:
    result: A tuple, containing (unit, min, max, avg, stdev).
    test: Name of test.
    iodepth: Iodepth parameter used by fio command.
  Returns:
    A list of samples in the form of 3 or 4 tuples. The tuples contain
        the sample metric (string), value (float), and unit (string).
        If a 4th element is included, it is a dictionary of sample
        metadata.
  """
  return [test, float(result[3]), result[0],
          {'min': result[1] + result[0],
           'max': result[2] + result[0],
           'stdev': result[4] + result[0],
           'max_jobs': FLAGS.maxjobs,
           'iodepth': iodepth}]

def RunSimulatedLogging(vm):
  """Spawn fio to simulate logging and gather the results.
  Args:
    vm: The vm that synthetic_storage_workloads_benchmark will be run upon.
  Returns:
    A list of samples in the form of 3 or 4 tuples. The tuples contain
        the sample metric (string), value (float), and unit (string).
        If a 4th element is included, it is a dictionary of sample
        metadata.
  """
  test_size = vm.total_memory_kb
  cmd = (
      'fio '
      '--filesize=10g '
      '--directory=%s '
      '--ioengine=libaio '
      '--filename=fio_test_file '
      '--invalidate=1 '
      '--randrepeat=0 '
      '--direct=0 '
      '--size=%dk '
      '--iodepth=%d ') % (vm.GetScratchDir(),
                          test_size,
                          DEFAULT_IODEPTH)
  if FLAGS.maxjobs:
    cmd += '--max-jobs=%s ' % FLAGS.maxjobs
  cmd += (
      '--name=sequential_write '
      '--overwrite=0 '
      '--rw=write '
      '--end_fsync=1 '
      '--name=random_read '
      '--size=%dk '
      '--stonewall '
      '--rw=randread '
      '--name=sequential_read '
      '--stonewall '
      '--rw=read ') % (test_size / 10)
  logging.info('FIO Results for simulated %s', LOGGING)
  res, _ = vm.RemoteCommand(cmd, should_log=True)
  latency, bandwidth = ParseFioResult(res)
  results = [
      CreateSampleFromBandwidthTuple(bandwidth[0],
                                     'sequential_write:bandwidth'),
      CreateSampleFromBandwidthTuple(bandwidth[1], 'random_read:bandwidth'),
      CreateSampleFromBandwidthTuple(bandwidth[2], 'sequential_read:bandwidth'),
      CreateSampleFromLatencyTuple(latency[0], 'sequential_write:latency'),
      CreateSampleFromLatencyTuple(latency[1], 'random_read:latency'),
      CreateSampleFromLatencyTuple(latency[2], 'sequential_read:latency')]
  return results

def RunSimulatedDatabase(vm):
  """Spawn fio to simulate database and gather the results.

  Args:
    vm: The vm that synthetic_storage_workloads_benchmark will be run upon.
  Returns:
    A list of samples in the form of 3 or 4 tuples. The tuples contain
        the same metric (string), value (float), and unit (string).
        If a 4th element is included, it is a dictionary of sample
        metadata.
  """
  test_size = min(vm.total_memory_kb / 10, 1000000)
  iodepth_list = FLAGS.iodepth_list or DEFAULT_DATABASE_SIMULATION_IODEPTH_LIST
  results = []
  for depth in iodepth_list:
    cmd = (
        'fio '
        '--filesize=10g '
        '--directory=%s '
        '--ioengine=libaio '
        '--filename=fio_test_file '
        '--overwrite=1 '
        '--invalidate=0 '
        '--direct=1 '
        '--randrepeat=0 '
        '--iodepth=%s '
        '--size=%dk '
        '--blocksize=4k ') % (vm.GetScratchDir(),
                              depth,
                              test_size)
    if FLAGS.maxjobs:
      cmd += '--max-jobs=%s ' % FLAGS.maxjobs
    cmd += (
        '--name=random_write '
        '--rw=randwrite '
        '--end_fsync=1 '
        '--name=random_read '
        '--stonewall '
        '--rw=randread '
        '--name=mixed_randrw '
        '--stonewall '
        '--rw=randrw '
        '--rwmixread=90 '
        '--rwmixwrite=10 '
        '--end_fsync=1 ')
    logging.info('FIO Results for simulated %s, iodepth %s', DATABASE, depth)
    res, _ = vm.RemoteCommand(cmd, should_log=True)
    latency, bandwidth = ParseFioResult(res)
    result = [CreateSampleFromBandwidthTuple(bandwidth[0],
                                             'random_write:bandwidth', depth),
              CreateSampleFromBandwidthTuple(bandwidth[1],
                                             'random_read:bandwidth', depth),
              CreateSampleFromBandwidthTuple(bandwidth[2],
                                             'mixed_randrw:read:bandwidth',
                                             depth),
              CreateSampleFromBandwidthTuple(bandwidth[3],
                                             'mixed_randrw:write:bandwidth',
                                             depth),
              CreateSampleFromLatencyTuple(latency[0], 'random_write:latency',
                                           depth),
              CreateSampleFromLatencyTuple(latency[1], 'random_read:latency',
                                           depth),
              CreateSampleFromLatencyTuple(latency[2],
                                           'mixed_randrw:read:latency', depth),
              CreateSampleFromLatencyTuple(latency[3],
                                           'mixed_randrw:write:latency', depth)]
    results.extend(result)
  return results

def RunSimulatedStreaming(vm):
  """Spawn fio to simulate streaming and gather the results.

  Args:
    vm: The vm that synthetic_storage_workloads_benchmark will be run upon.
  Returns:
    A list of samples in the form of 3 or 4 tuples. The tuples contain
        the same metric (string), value (float), and unit (string).
        If a 4th element is included, it is a dictionary of sample
        metadata.
  """
  test_size = min(vm.total_memory_kb / 10, 1000000)
  iodepth_list = FLAGS.iodepth_list or DEFAULT_STREAMING_SIMULATION_IODEPTH_LIST
  results = []
  for depth in iodepth_list:
    cmd = (
        'fio '
        '--filesize=10g '
        '--directory=%s '
        '--ioengine=libaio '
        '--overwrite=0 '
        '--invalidate=1 '
        '--direct=1 '
        '--randrepeat=0 '
        '--iodepth=%s '
        '--blocksize=1m '
        '--size=%dk '
        '--filename=fio_test_file ') % (vm.GetScratchDir(),
                                        depth,
                                        test_size)
    if FLAGS.maxjobs:
      cmd += '--max-jobs=%s ' % FLAGS.maxjobs
    cmd += (
        '--name=sequential_write '
        '--rw=write '
        '--end_fsync=1 '
        '--name=sequential_read '
        '--stonewall '
        '--rw=read ')
    logging.info('FIO Results for simulated %s', STREAMING)
    res, _ = vm.RemoteCommand(cmd, should_log=True)
    latency, bandwidth = ParseFioResult(res)
    result = [
        CreateSampleFromBandwidthTuple(bandwidth[0],
                                       'sequential_write:bandwidth', depth),
        CreateSampleFromBandwidthTuple(bandwidth[1],
                                       'sequential_read:bandwidth', depth),
        CreateSampleFromLatencyTuple(latency[0],
                                     'sequential_write:latency', depth),
        CreateSampleFromLatencyTuple(latency[1],
                                     'sequential_read:latency', depth)]
    results.extend(result)
  return results

RUN_SCENARIO_FUNCTION_DICT = {
    LOGGING: {DESCRIPTION: 'simulated_logging', METHOD: RunSimulatedLogging},
    DATABASE: {DESCRIPTION: 'simulated_database', METHOD: RunSimulatedDatabase},
    STREAMING: {DESCRIPTION: 'simulated_streaming',
                METHOD: RunSimulatedStreaming}}

def Run(benchmark_spec):
  """Spawn fio and gather the results.

  Args:
    benchmark_spec: The benchmark specification. Contains all data that is
        required to run the benchmark.

  Returns:
    A list of samples in the form of 3 or 4 tuples. The tuples contain
        the sample metric (string), value (float), and unit (string).
        If a 4th element is included, it is a dictionary of sample
        metadata.
  """
  logging.info('Simulating %s senario.', FLAGS.workload_mode)
  vms = benchmark_spec.vms
  vm = vms[0]
  # Add mode name into benchmark name, so perfkitbenchmarker will publish each
  # mode as a different benchmark, instead of mixing them together.
  BENCHMARK_INFO['name'] = '%s:%s' % (BENCHMARK_INFO['name'].split(':')[0],
                                      FLAGS.workload_mode)
  return RUN_SCENARIO_FUNCTION_DICT[FLAGS.workload_mode][METHOD](vm)

def Cleanup(benchmark_spec):
  """Uninstall packages required for fio and remove benchmark files.

  Args:
    benchmark_spec: The benchmark specification. Contains all data that is
        required to run the benchmark.
  """
  vms = benchmark_spec.vms
  vm = vms[0]
  logging.info('FIO Cleanup up on %s', vm)
  vm.UninstallPackage(REQUIRED_PACKAGES)
  vm.RemoveFile(vm.GetScratchDir() + '/fio_test_file')
<commit_msg>Get rid of queue depth equal to 1 benchmark for database scenarios.<commit_after># Copyright 2014 Google Inc. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Runs fio benchmarks to simulate logging, database and streaming.

Man: http://manpages.ubuntu.com/manpages/natty/man1/fio.1.html
Quick howto: http://www.bluestop.org/fio/HOWTO.txt

Simulated logging benchmark does the following things (PD only):
0) Do NOT use direct IO for any tests below, simply go through the FS.
1) Sequentially write x GB with queue depth equal to 8, where x is decided by
   the test VM's total memory. (A larger VM will write more bytes)
2) Random read of 10% of the bytes written.
3) Sequential read of all of the bytes written.

Simulated database benchmark does the following things (PD, PD-SSD, local SSD):
1) 4K Random R on a file using queue depths 1, 16 and 64 (each queue depth
   is a different benchmark).
2) 4K Random W on a file using queue depths 1, 16 and 64 (each queue depth
   is a different benchmark).
3) 4K Random 90% R/ 10% W on a file using queue depths 1, 16 and 64 (each
   queue depth is a different benchmark).
4) The size of the test file is decided by the test VM's total memory and capped
   at 1GB to ensure this test finishes within reasonable time.

Simulated streaming benchmark (PD only):
1) 1M Seq R at queue depth 1 and 16 (streaming).
2) 1M Seq W at queue depth 1 and 16 (streaming).

For AWS, where use PD, we should use EBS-GP and EBS Magnetic, for PD-SSD use
EBS-GP and PIOPS.
"""

import logging
import re

from perfkitbenchmarker import flags

LOGGING = 'logging'
DATABASE = 'database'
STREAMING = 'streaming'

flags.DEFINE_enum('workload_mode', LOGGING,
                  [LOGGING, DATABASE, STREAMING],
                  'Simulate a logging, database or streaming scenario.')

flags.DEFINE_list('iodepth_list', None, 'A list of iodepth parameter used by '
                  'fio command in simulated database and streaming scenarios '
                  'only.')

flags.DEFINE_integer('maxjobs', 0,
                     'The maximum allowed number of jobs to support.')

FLAGS = flags.FLAGS

BENCHMARK_INFO = {'name': 'synthetic_storage_workloads_benchmark',
                  'description': 'Runs FIO in sequenctial, random, read and '
                                 'write modes to simulate various scenarios.',
                  'scratch_disk': True,
                  'num_machines': 1}
REQUIRED_PACKAGES = 'bc fio libaio1'
DESCRIPTION = 'description'
METHOD = 'method'

DEFAULT_IODEPTH = 8
DEFAULT_DATABASE_SIMULATION_IODEPTH_LIST = [16, 64]
DEFAULT_STREAMING_SIMULATION_IODEPTH_LIST = [1, 16]

LATENCY_REGEX = r'[=\s]+([\d\.]+)[\s,]+'
BANDWIDTH_REGEX = r'(\d+)(\w+/*\w*)'

def GetInfo():
  return BENCHMARK_INFO

def Prepare(benchmark_spec):
  """Prepare the virtual machine to run FIO.

     This includes installing fio, bc. and libaio1 and insuring that the
     attached disk is large enough to support the fio benchmark.

  Args:
    benchmark_spec: The benchmark specification. Contains all data that is
        required to run the benchmark.
  """
  vms = benchmark_spec.vms
  vm = vms[0]
  logging.info('FIO prepare on %s', vm)
  vm.InstallPackage(REQUIRED_PACKAGES)

def ParseFioResult(res):
  """Parse result from fio commands.

  Args:
    res: Result from fio commands in string format.
  Returns:
    A list of tuples represents latency. Each tuple contains unit and
    latency number.
    A list of tuples represents bandwidth. Each tuple contains aggrb,
    unit of aggrb, minb, unit of minb, maxb, unit of maxb, mint, unit of mint,
    maxt, unit of maxt.
  """
  latency = re.findall(r'\s+lat \((\w+)\)[\s:]+'
                       r'min%smax%savg%sstdev%s' % (
                           LATENCY_REGEX, LATENCY_REGEX,
                           LATENCY_REGEX, LATENCY_REGEX), res)
  bandwidth = re.findall(
      r'aggrb=%s, minb=%s, maxb=%s, mint=%s, maxt=%s' % (
          BANDWIDTH_REGEX, BANDWIDTH_REGEX, BANDWIDTH_REGEX, BANDWIDTH_REGEX,
          BANDWIDTH_REGEX), res)
  print latency, bandwidth
  return latency, bandwidth

def CreateSampleFromBandwidthTuple(result, test, iodepth=DEFAULT_IODEPTH):
  """Create a sample from bandwidth result tuple.

  Args:
    result: A tuple, containing (aggrb, unit of aggrb, minb, unit of minb,
        maxb, unit of maxb, mint, unit of mint, maxt, unit of maxt).
    test: Name of test.
    iodepth: Iodepth parameter used by fio command.
  Returns:
    A list of samples in the form of 3 or 4 tuples. The tuples contain
        the sample metric (string), value (float), and unit (string).
        If a 4th element is included, it is a dictionary of sample
        metadata.
  """
  return [test, float(result[0]), result[1],
          {'minb': result[2] + result[3],
           'maxb': result[4] + result[5],
           'mint': result[6] + result[7],
           'maxt': result[8] + result[9],
           'max_jobs': FLAGS.maxjobs,
           'iodepth': iodepth}]

def CreateSampleFromLatencyTuple(result, test, iodepth=DEFAULT_IODEPTH):
  """Create a sample from latency result tuple.

  Args:
    result: A tuple, containing (unit, min, max, avg, stdev).
    test: Name of test.
    iodepth: Iodepth parameter used by fio command.
  Returns:
    A list of samples in the form of 3 or 4 tuples. The tuples contain
        the sample metric (string), value (float), and unit (string).
        If a 4th element is included, it is a dictionary of sample
        metadata.
  """
  return [test, float(result[3]), result[0],
          {'min': result[1] + result[0],
           'max': result[2] + result[0],
           'stdev': result[4] + result[0],
           'max_jobs': FLAGS.maxjobs,
           'iodepth': iodepth}]

def RunSimulatedLogging(vm):
  """Spawn fio to simulate logging and gather the results.
  Args:
    vm: The vm that synthetic_storage_workloads_benchmark will be run upon.
  Returns:
    A list of samples in the form of 3 or 4 tuples. The tuples contain
        the sample metric (string), value (float), and unit (string).
        If a 4th element is included, it is a dictionary of sample
        metadata.
  """
  test_size = vm.total_memory_kb
  cmd = (
      'fio '
      '--filesize=10g '
      '--directory=%s '
      '--ioengine=libaio '
      '--filename=fio_test_file '
      '--invalidate=1 '
      '--randrepeat=0 '
      '--direct=0 '
      '--size=%dk '
      '--iodepth=%d ') % (vm.GetScratchDir(),
                          test_size,
                          DEFAULT_IODEPTH)
  if FLAGS.maxjobs:
    cmd += '--max-jobs=%s ' % FLAGS.maxjobs
  cmd += (
      '--name=sequential_write '
      '--overwrite=0 '
      '--rw=write '
      '--end_fsync=1 '
      '--name=random_read '
      '--size=%dk '
      '--stonewall '
      '--rw=randread '
      '--name=sequential_read '
      '--stonewall '
      '--rw=read ') % (test_size / 10)
  logging.info('FIO Results for simulated %s', LOGGING)
  res, _ = vm.RemoteCommand(cmd, should_log=True)
  latency, bandwidth = ParseFioResult(res)
  results = [
      CreateSampleFromBandwidthTuple(bandwidth[0],
                                     'sequential_write:bandwidth'),
      CreateSampleFromBandwidthTuple(bandwidth[1], 'random_read:bandwidth'),
      CreateSampleFromBandwidthTuple(bandwidth[2], 'sequential_read:bandwidth'),
      CreateSampleFromLatencyTuple(latency[0], 'sequential_write:latency'),
      CreateSampleFromLatencyTuple(latency[1], 'random_read:latency'),
      CreateSampleFromLatencyTuple(latency[2], 'sequential_read:latency')]
  return results

def RunSimulatedDatabase(vm):
  """Spawn fio to simulate database and gather the results.

  Args:
    vm: The vm that synthetic_storage_workloads_benchmark will be run upon.
  Returns:
    A list of samples in the form of 3 or 4 tuples. The tuples contain
        the same metric (string), value (float), and unit (string).
        If a 4th element is included, it is a dictionary of sample
        metadata.
  """
  test_size = min(vm.total_memory_kb / 10, 1000000)
  iodepth_list = FLAGS.iodepth_list or DEFAULT_DATABASE_SIMULATION_IODEPTH_LIST
  results = []
  for depth in iodepth_list:
    cmd = (
        'fio '
        '--filesize=10g '
        '--directory=%s '
        '--ioengine=libaio '
        '--filename=fio_test_file '
        '--overwrite=1 '
        '--invalidate=0 '
        '--direct=1 '
        '--randrepeat=0 '
        '--iodepth=%s '
        '--size=%dk '
        '--blocksize=4k ') % (vm.GetScratchDir(),
                              depth,
                              test_size)
    if FLAGS.maxjobs:
      cmd += '--max-jobs=%s ' % FLAGS.maxjobs
    cmd += (
        '--name=random_write '
        '--rw=randwrite '
        '--end_fsync=1 '
        '--name=random_read '
        '--stonewall '
        '--rw=randread '
        '--name=mixed_randrw '
        '--stonewall '
        '--rw=randrw '
        '--rwmixread=90 '
        '--rwmixwrite=10 '
        '--end_fsync=1 ')
    logging.info('FIO Results for simulated %s, iodepth %s', DATABASE, depth)
    res, _ = vm.RemoteCommand(cmd, should_log=True)
    latency, bandwidth = ParseFioResult(res)
    result = [CreateSampleFromBandwidthTuple(bandwidth[0],
                                             'random_write:bandwidth', depth),
              CreateSampleFromBandwidthTuple(bandwidth[1],
                                             'random_read:bandwidth', depth),
              CreateSampleFromBandwidthTuple(bandwidth[2],
                                             'mixed_randrw:read:bandwidth',
                                             depth),
              CreateSampleFromBandwidthTuple(bandwidth[3],
                                             'mixed_randrw:write:bandwidth',
                                             depth),
              CreateSampleFromLatencyTuple(latency[0], 'random_write:latency',
                                           depth),
              CreateSampleFromLatencyTuple(latency[1], 'random_read:latency',
                                           depth),
              CreateSampleFromLatencyTuple(latency[2],
                                           'mixed_randrw:read:latency', depth),
              CreateSampleFromLatencyTuple(latency[3],
                                           'mixed_randrw:write:latency', depth)]
    results.extend(result)
  return results

def RunSimulatedStreaming(vm):
  """Spawn fio to simulate streaming and gather the results.

  Args:
    vm: The vm that synthetic_storage_workloads_benchmark will be run upon.
  Returns:
    A list of samples in the form of 3 or 4 tuples. The tuples contain
        the same metric (string), value (float), and unit (string).
        If a 4th element is included, it is a dictionary of sample
        metadata.
  """
  test_size = min(vm.total_memory_kb / 10, 1000000)
  iodepth_list = FLAGS.iodepth_list or DEFAULT_STREAMING_SIMULATION_IODEPTH_LIST
  results = []
  for depth in iodepth_list:
    cmd = (
        'fio '
        '--filesize=10g '
        '--directory=%s '
        '--ioengine=libaio '
        '--overwrite=0 '
        '--invalidate=1 '
        '--direct=1 '
        '--randrepeat=0 '
        '--iodepth=%s '
        '--blocksize=1m '
        '--size=%dk '
        '--filename=fio_test_file ') % (vm.GetScratchDir(),
                                        depth,
                                        test_size)
    if FLAGS.maxjobs:
      cmd += '--max-jobs=%s ' % FLAGS.maxjobs
    cmd += (
        '--name=sequential_write '
        '--rw=write '
        '--end_fsync=1 '
        '--name=sequential_read '
        '--stonewall '
        '--rw=read ')
    logging.info('FIO Results for simulated %s', STREAMING)
    res, _ = vm.RemoteCommand(cmd, should_log=True)
    latency, bandwidth = ParseFioResult(res)
    result = [
        CreateSampleFromBandwidthTuple(bandwidth[0],
                                       'sequential_write:bandwidth', depth),
        CreateSampleFromBandwidthTuple(bandwidth[1],
                                       'sequential_read:bandwidth', depth),
        CreateSampleFromLatencyTuple(latency[0],
                                     'sequential_write:latency', depth),
        CreateSampleFromLatencyTuple(latency[1],
                                     'sequential_read:latency', depth)]
    results.extend(result)
  return results

RUN_SCENARIO_FUNCTION_DICT = {
    LOGGING: {DESCRIPTION: 'simulated_logging', METHOD: RunSimulatedLogging},
    DATABASE: {DESCRIPTION: 'simulated_database', METHOD: RunSimulatedDatabase},
    STREAMING: {DESCRIPTION: 'simulated_streaming',
                METHOD: RunSimulatedStreaming}}

def Run(benchmark_spec):
  """Spawn fio and gather the results.

  Args:
    benchmark_spec: The benchmark specification. Contains all data that is
        required to run the benchmark.

  Returns:
    A list of samples in the form of 3 or 4 tuples. The tuples contain
        the sample metric (string), value (float), and unit (string).
        If a 4th element is included, it is a dictionary of sample
        metadata.
  """
  logging.info('Simulating %s senario.', FLAGS.workload_mode)
  vms = benchmark_spec.vms
  vm = vms[0]
  # Add mode name into benchmark name, so perfkitbenchmarker will publish each
  # mode as a different benchmark, instead of mixing them together.
  BENCHMARK_INFO['name'] = '%s:%s' % (BENCHMARK_INFO['name'].split(':')[0],
                                      FLAGS.workload_mode)
  return RUN_SCENARIO_FUNCTION_DICT[FLAGS.workload_mode][METHOD](vm)

def Cleanup(benchmark_spec):
  """Uninstall packages required for fio and remove benchmark files.

  Args:
    benchmark_spec: The benchmark specification. Contains all data that is
        required to run the benchmark.
  """
  vms = benchmark_spec.vms
  vm = vms[0]
  logging.info('FIO Cleanup up on %s', vm)
  vm.UninstallPackage(REQUIRED_PACKAGES)
  vm.RemoveFile(vm.GetScratchDir() + '/fio_test_file')
<|endoftext|>