<commit_before>###############################################################################
#   actionAngle: a Python module to calculate  actions, angles, and frequencies
#
#      class: actionAnglePower
#
#      methods:
#              JR
#              Jphi
#              angleR
#              TR
#              Tphi
#              I
#              calcRapRperi
#              potential
#              calcEL
###############################################################################
import math as m
import numpy as nu
from scipy import optimize, integrate
from actionAngle import *
from actionAngleVertical import actionAngleVertical
class actionAnglePower(actionAngle,actionAngleVertical):
    """Action-angle formalism for power-law rotation curves"""
    def __init__(self,*args,**kwargs):
        """
        NAME:
           __init__
        PURPOSE:
           initialize an actionAnglePower object, power-law is vc=vo(r)^beta
        INPUT:
           Either:
              a) R,vR,vT
              b) Orbit instance: initial condition used if that's it, orbit(t)
                 if there is a time given as well
           beta= power-law index
        OUTPUT:
        HISTORY:
           2010-07-11 - Written - Bovy (NYU)
        """
        actionAngle.__init__(self,*args,**kwargs)
        if kwargs.has_key('verticalPot'):
            actionAngleVertical.__init__(self,*args,pot=kwargs['verticalPot'],
                                         **kwargs)
        if not kwargs.has_key('beta'):
            raise IOError("Must specify beta= for actionAnglePower")
        self._beta= kwargs['beta']
        if self._beta < 0.:
            self._signbeta= -1
        else:
            self._signbeta= 1
        self._absbeta= nu.fabs(self._beta)
        (rperi,rap)= self.calcRapRperi()
        self._e= (rap-rperi)/(rap+rperi)
        if self._e == 0.:
            self._X= 0.
            self._Y= 0.
        else:
            self._X= (1.-self._e)**2./4./self._e\
                *(((1.+self._e)/(1.-self._e))**(2.*self._beta+2.)-1.)
            self._Y= (1.+self._e)**2./4./self._e\
                *(1.-((1.-self._e)/(1.+self._e))**(2.*self._beta+2.))
        return None

    def angleR(self,**kwargs):
        """
        NAME:
           AngleR
        PURPOSE:
           Calculate the radial angle
        INPUT:
           scipy.integrate.quadrature keywords
        OUTPUT:
           w_R(R,vT,vT) in radians +
           estimate of the error (does not include TR error)
        HISTORY:
           2010-11-30 - Written - Bovy (NYU)
        """
        if hasattr(self,'_angleR'):
            return self._angleR
        (rperi,rap)= self.calcRapRperi()
        if rap == rperi:
            return 0.
        TR= self.TR(**kwargs)[0]
        Rmean= m.exp((m.log(rperi)+m.log(rap))/2.)
        if self._R < Rmean:
            if self._R > rperi:
                wR= (2.*m.pi/TR*m.sqrt(self._absbeta)*rperi**(1.-self._beta)*
                     nu.array(integrate.quadrature(_TRPowerIntegrandSmall,
                                                   0.,m.sqrt(self._R/rperi-1.),
                                                   args=(self._X,self._beta,
                                                         self._signbeta),
                                                   **kwargs)))\
                                                   +nu.array([m.pi,0.])
            else:
                wR= nu.array([m.pi,0.])
        else:
            if self._R < rap:
                wR= -(2.*m.pi/TR*m.sqrt(self._absbeta)*rap**(1.-self._beta)*
                      nu.array(integrate.quadrature(_TRPowerIntegrandLarge,
                                                    0.,m.sqrt(1.-self._R/rap),
                                                    args=(self._Y,self._beta,
                                                          self._signbeta),
                                                    **kwargs)))
            else:
                wR= nu.array([0.,0.])
        if self._vR < 0.:
            wR[0]+= m.pi
        self._angleR= nu.array([wR[0] % (2.*m.pi),wR[1]])
        return self._angleR

    def TR(self,**kwargs):
        """
        NAME:
           TR
        PURPOSE:
           Calculate the radial period for a power-law rotation curve
        INPUT:
           scipy.integrate.quadrature keywords
        OUTPUT:
           T_R(R,vT,vT)*vc/ro + estimate of the error
        HISTORY:
           2010-11-30 - Written - Bovy (NYU)
        """
        if hasattr(self,'_TR'):
            return self._TR
        (rperi,rap)= self.calcRapRperi()
        if rap == rperi: #Rough limit
            #TR=kappa
            gamma= m.sqrt(2./(1.+self._beta))
            kappa= 2.*self._R**(self._beta-1.)/gamma
            self._TR= nu.array([2.*m.pi/kappa,0.])
            return self._TR
        Rmean= m.exp((m.log(rperi)+m.log(rap))/2.)
        TR= 0.
        if Rmean > rperi:
            TR+= rperi**(1.-self._beta)\
                *nu.array(integrate.quadrature(_TRPowerIntegrandSmall,
                                               0.,m.sqrt(Rmean/rperi-1.),
                                               args=(self._X,self._beta,
                                                     self._signbeta),
                                               **kwargs))
        if Rmean < rap:
            TR+= rap**(1.-self._beta)\
                *nu.array(integrate.quadrature(_TRPowerIntegrandLarge,
                                               0.,m.sqrt(1.-Rmean/rap),
                                               args=(self._Y,self._beta,
                                                     self._signbeta),
                                               **kwargs))
        self._TR= 2.*m.sqrt(self._absbeta)*TR
        return self._TR

    def Tphi(self,**kwargs):
        """
        NAME:
           Tphi
        PURPOSE:
           Calculate the azimuthal period for a power-law rotation curve
        INPUT:
           +scipy.integrate.quadrature keywords
        OUTPUT:
           T_phi(R,vT,vT)/ro/vc + estimate of the error
        HISTORY:
           2010-11-30 - Written - Bovy (NYU)
        """
        if hasattr(self,'_Tphi'):
            return self._Tphi
        (rperi,rap)= self.calcRapRperi()
        if rap == rperi:#Circular orbit
            return nu.array([2.*m.pi*self._R/self._vT,0.])
        TR= self.TR(**kwargs)
        I= self.I(**kwargs)
        Tphi= nu.zeros(2)
        Tphi[0]= TR[0]/I[0]*m.pi
        Tphi[1]= Tphi[0]*m.sqrt((I[1]/I[0])**2.+(TR[1]/TR[0])**2.)
        self._Tphi= Tphi
        return self._Tphi

    def I(self,**kwargs):
        """
        NAME:
           I
        PURPOSE:
           Calculate I, the 'ratio' between the radial and azimutha period,
           for a power-law rotation curve
        INPUT:
           +scipy.integrate.quadrature keywords
        OUTPUT:
           I(R,vT,vT) + estimate of the error
        HISTORY:
           2010-11-30 - Written - Bovy (NYU)
        """
        if hasattr(self,'_I'):
            return self._I
        (rperi,rap)= self.calcRapRperi()
        Rmean= m.exp((m.log(rperi)+m.log(rap))/2.)
        if rap == rperi: #Rough limit
            TR= self.TR()[0]
            Tphi= self.Tphi()[0]
            self._I= nu.array([TR/Tphi,0.])
            return self._I
        I= 0.
        if Rmean > rperi:
            I+= rperi**(-1.-self._beta)\
                *nu.array(integrate.quadrature(_IPowerIntegrandSmall,
                                               0.,m.sqrt(Rmean/rperi-1.),
                                               args=(self._X,self._beta,
                                                     self._signbeta),
                                               **kwargs))
        if Rmean < rap:
            I+= rap**(-1.-self._beta)\
                *nu.array(integrate.quadrature(_IPowerIntegrandLarge,
                                               0.,m.sqrt(1.-Rmean/rap),
                                               args=(self._Y,self._beta,
                                                     self._signbeta),
                                               **kwargs))
        self._I= I*self._R*self._vT*m.sqrt(self._absbeta)
        return self._I

    def Jphi(self):
        """
        NAME:
           Jphi
        PURPOSE:
           Calculate the azimuthal action
        INPUT:
        OUTPUT:
           J_R(R,vT,vT)/ro/vc
        HISTORY:
           2010-05-13 - Written - Bovy (NYU)
        """
        return nu.array([self._R*self._vT,0.])

    def JR(self,**kwargs):
        """
        NAME:
           JR
        PURPOSE:
           Calculate the radial action for a power-law rotation curve
        INPUT:
           +scipy.integrate.quad keywords
        OUTPUT:
           J_R(R,vT,vT)/ro/vc + estimate of the error
        HISTORY:
           2010-11-30 - Written - Bovy (NYU)
        """
        if hasattr(self,'_JR'):
            return self._JR
        (rperi,rap)= self.calcRapRperi()
        self._JR= (2./m.sqrt(self._absbeta)*rperi**(self._beta+1.)*
                   nu.array(integrate.quad(_JRPowerIntegrand,1.,rap/rperi,
                                           args=(self._X,self._beta,
                                                 self._signbeta),
                                           **kwargs)))
        return self._JR

    def calcRapRperi(self):
        """
        NAME:
           calcRapRperi
        PURPOSE:
           calculate the apocenter and pericenter radii for a power-law
           rotation curve
        INPUT:
        OUTPUT:
           (rperi,rap)
        HISTORY:
           2010-11-30 - Written - Bovy (NYU)
        """
        if hasattr(self,'_rperirap'):
            return self._rperirap
        EL= calcELPower(self._R,self._vR,self._vT,self._beta,vc=1.,ro=1.)
        E, L= EL
        if self._vR == 0. and self._vT > self._R**self._beta: #We are exactly at pericenter
            rperi= self._R
            rend= _rapRperiPowerFindStart(self._R,E,L,self._beta,rap=True)
            rap= optimize.brentq(_rapRperiPowerEq,rperi+0.000001,
                                 rend,args=(E,L,self._beta))
#                                 fprime=_rapRperiPowerDeriv)
        elif self._vR == 0. and self._vT <self.R**self._beta: #We are exactly at apocenter
            rap= self._R
            rstart= _rapRperiPowerFindStart(self._R,E,L,self._beta)
            rperi= optimize.brentq(_rapRperiPowerEq,rstart,rap-0.000001,
                                   args=(E,L,self._beta))
#                                   fprime=_rapRperiPowerDeriv)
        elif self._vR == 0. and self._vT == self._R**self.__beta: #We are on a circular orbit
            rperi= self._R
            rap = self._R
        else:
            rstart= _rapRperiPowerFindStart(self._R,E,L,self._beta)
            rperi= optimize.brentq(_rapRperiPowerEq,rstart,self._R,
                                   (E,L,self._beta))
            rend= _rapRperiPowerFindStart(self._R,E,L,self._beta,rap=True)
            rap= optimize.brentq(_rapRperiPowerEq,self._R,rend,
                                 (E,L,self._beta))
        self._rperirap= (rperi,rap)
        return self._rperirap

def calcRapRperiFromELPower(E,L,beta,vc=1.,ro=1.):
    """
    NAME:
       calcRapRperiFromELPower
    PURPOSE:
       calculate the apocenter and pericenter radii for a power-law
       rotation curve
    INPUT:
       E - energy
       L - angular momemtum
       beta - power-law index of rotation curve
       vc - circular velocity
       ro - reference radius
    OUTPUT:
       (rperi,rap)
    HISTORY:
       2010-11-30 - Written - Bovy (NYU)
    """
    rstart= _rapRperiPowerFindStart(L,E,L)
    rperi= optimize.brentq(_rapRperiPowerEq,rstart,L,(E,L,beta))
    rend= _rapRperiPowerFindStart(L,E,L,rap=True)
    rap= optimize.brentq(_rapRperiPowerEq,L,rend,(E,L,beta))
    return (rperi,rap)

def calcELPower(R,vR,vT,beta,vc=1.,ro=1.):
    """
    NAME:
       calcELPower
    PURPOSE:
       calculate the energy and angular momentum for a flat rotation curve
    INPUT:
       R - Galactocentric radius (/ro)
       vR - radial part of the velocity (/vc)
       vT - azimuthal part of the velocity (/vc)
       vc - circular velocity
       ro - reference radius
    OUTPUT:
       (E,L)
    HISTORY:
       2010-11-30 - Written - Bovy (NYU)
    """
    return (potentialPower(R,beta,vc=vc,ro=ro)+vR**2./2.+vT**2./2.,R*vT)

def potentialPower(R,beta,vc=1.,ro=1.):
    """
    NAME:
       potentialPower
    PURPOSE:
       return the potential for a power-law rotation curve
    INPUT:
       R - Galactocentric radius (/ro)
       beta - power-law index
       vc - circular velocity
       ro - reference radius
    OUTPUT:
       Phi(R)
    HISTORY:
       2010-11-30 - Written - Bovy (NYU)
    """
    return vc**2./2./beta*(R/ro)**(2.*beta)

def _JRPowerIntegrand(r,X,beta,signbeta):
    """The J_R integrand for a flat rotation curve"""
    return nu.sqrt(signbeta*(1.-X-r**(2.*beta+2.)+r**2.*X))/r

def _TRPowerIntegrandSmall(t,X,beta,signbeta):
    r= 1.+t**2.#part of the transformation
    return 2.*t/_JRPowerIntegrand(r,X,beta,signbeta)

def _TRPowerIntegrandLarge(t,Y,beta,signbeta):
    r= 1.-t**2.#part of the transformation
    return 2.*t/_JRPowerIntegrand(r,Y,beta,signbeta) #same integrand

def _IPowerIntegrandSmall(t,X,beta,signbeta):
    r= 1.+t**2.#part of the transformation
    return 2.*t/_JRPowerIntegrand(r,X,beta,signbeta)/r**2.

def _IPowerIntegrandLarge(t,Y,beta,signbeta):
    r= 1.-t**2.#part of the transformation
    return 2.*t/_JRPowerIntegrand(r,Y,beta,signbeta)/r**2.

def _rapRperiPowerEq(R,E,L,beta):
    """The vr=0 equation that needs to be solved to find apo- and pericenter"""
    return E-potentialPower(R,beta)-L**2./2./R**2.

def _rapRperiPowerDeriv(R,E,L,beta):
    """The derivative of the vr=0 equation that needs to be solved to find
    apo- and pericenter"""
    return -R**(2.*beta-1.)+L**2./R**3.

def _rapRperiPowerFindStart(R,E,L,beta,rap=False):
    """
    NAME:
       _rapRperiPowerFindStart
    PURPOSE:
       Find adequate start or end points to solve for rap and rperi
    INPUT:
       R - Galactocentric radius
       E - energy
       L - angular momentum
       beta - power-law index of rotation curve
       rap - if True, find the rap end-point
    OUTPUT:
       rstart or rend
    HISTORY:
       2010-11-30 - Written - Bovy (NYU)
    """
    if rap:
        rtry= 2.*R
    else:
        rtry= R/2.
    while (E-potentialPower(rtry,beta)-L**2./2./rtry**2) > 0.:
        if rap:
            rtry*= 2.
        else:
            rtry/= 2.
    return rtry

<commit_msg>add UnboundError to actionAngle<commit_after>###############################################################################
#   actionAngle: a Python module to calculate  actions, angles, and frequencies
#
#      class: actionAnglePower
#
#      methods:
#              JR
#              Jphi
#              angleR
#              TR
#              Tphi
#              I
#              calcRapRperi
#              potential
#              calcEL
###############################################################################
import math as m
import numpy as nu
from scipy import optimize, integrate
from actionAngle import *
from actionAngleVertical import actionAngleVertical
class actionAnglePower(actionAngle,actionAngleVertical):
    """Action-angle formalism for power-law rotation curves"""
    def __init__(self,*args,**kwargs):
        """
        NAME:
           __init__
        PURPOSE:
           initialize an actionAnglePower object, power-law is vc=vo(r)^beta
        INPUT:
           Either:
              a) R,vR,vT
              b) Orbit instance: initial condition used if that's it, orbit(t)
                 if there is a time given as well
           beta= power-law index
        OUTPUT:
        HISTORY:
           2010-07-11 - Written - Bovy (NYU)
        """
        actionAngle.__init__(self,*args,**kwargs)
        if kwargs.has_key('verticalPot'):
            actionAngleVertical.__init__(self,*args,pot=kwargs['verticalPot'],
                                         **kwargs)
        if not kwargs.has_key('beta'):
            raise IOError("Must specify beta= for actionAnglePower")
        self._beta= kwargs['beta']
        if self._beta < 0.:
            self._signbeta= -1
        else:
            self._signbeta= 1
        self._absbeta= nu.fabs(self._beta)
        (rperi,rap)= self.calcRapRperi()
        self._e= (rap-rperi)/(rap+rperi)
        if self._e == 0.:
            self._X= 0.
            self._Y= 0.
        else:
            self._X= (1.-self._e)**2./4./self._e\
                *(((1.+self._e)/(1.-self._e))**(2.*self._beta+2.)-1.)
            self._Y= (1.+self._e)**2./4./self._e\
                *(1.-((1.-self._e)/(1.+self._e))**(2.*self._beta+2.))
        return None

    def angleR(self,**kwargs):
        """
        NAME:
           AngleR
        PURPOSE:
           Calculate the radial angle
        INPUT:
           scipy.integrate.quadrature keywords
        OUTPUT:
           w_R(R,vT,vT) in radians +
           estimate of the error (does not include TR error)
        HISTORY:
           2010-11-30 - Written - Bovy (NYU)
        """
        if hasattr(self,'_angleR'):
            return self._angleR
        (rperi,rap)= self.calcRapRperi()
        if rap == rperi:
            return 0.
        TR= self.TR(**kwargs)[0]
        Rmean= m.exp((m.log(rperi)+m.log(rap))/2.)
        if self._R < Rmean:
            if self._R > rperi:
                wR= (2.*m.pi/TR*m.sqrt(self._absbeta)*rperi**(1.-self._beta)*
                     nu.array(integrate.quadrature(_TRPowerIntegrandSmall,
                                                   0.,m.sqrt(self._R/rperi-1.),
                                                   args=(self._X,self._beta,
                                                         self._signbeta),
                                                   **kwargs)))\
                                                   +nu.array([m.pi,0.])
            else:
                wR= nu.array([m.pi,0.])
        else:
            if self._R < rap:
                wR= -(2.*m.pi/TR*m.sqrt(self._absbeta)*rap**(1.-self._beta)*
                      nu.array(integrate.quadrature(_TRPowerIntegrandLarge,
                                                    0.,m.sqrt(1.-self._R/rap),
                                                    args=(self._Y,self._beta,
                                                          self._signbeta),
                                                    **kwargs)))
            else:
                wR= nu.array([0.,0.])
        if self._vR < 0.:
            wR[0]+= m.pi
        self._angleR= nu.array([wR[0] % (2.*m.pi),wR[1]])
        return self._angleR

    def TR(self,**kwargs):
        """
        NAME:
           TR
        PURPOSE:
           Calculate the radial period for a power-law rotation curve
        INPUT:
           scipy.integrate.quadrature keywords
        OUTPUT:
           T_R(R,vT,vT)*vc/ro + estimate of the error
        HISTORY:
           2010-11-30 - Written - Bovy (NYU)
        """
        if hasattr(self,'_TR'):
            return self._TR
        (rperi,rap)= self.calcRapRperi()
        if rap == rperi: #Rough limit
            #TR=kappa
            gamma= m.sqrt(2./(1.+self._beta))
            kappa= 2.*self._R**(self._beta-1.)/gamma
            self._TR= nu.array([2.*m.pi/kappa,0.])
            return self._TR
        Rmean= m.exp((m.log(rperi)+m.log(rap))/2.)
        TR= 0.
        if Rmean > rperi:
            TR+= rperi**(1.-self._beta)\
                *nu.array(integrate.quadrature(_TRPowerIntegrandSmall,
                                               0.,m.sqrt(Rmean/rperi-1.),
                                               args=(self._X,self._beta,
                                                     self._signbeta),
                                               **kwargs))
        if Rmean < rap:
            TR+= rap**(1.-self._beta)\
                *nu.array(integrate.quadrature(_TRPowerIntegrandLarge,
                                               0.,m.sqrt(1.-Rmean/rap),
                                               args=(self._Y,self._beta,
                                                     self._signbeta),
                                               **kwargs))
        self._TR= 2.*m.sqrt(self._absbeta)*TR
        return self._TR

    def Tphi(self,**kwargs):
        """
        NAME:
           Tphi
        PURPOSE:
           Calculate the azimuthal period for a power-law rotation curve
        INPUT:
           +scipy.integrate.quadrature keywords
        OUTPUT:
           T_phi(R,vT,vT)/ro/vc + estimate of the error
        HISTORY:
           2010-11-30 - Written - Bovy (NYU)
        """
        if hasattr(self,'_Tphi'):
            return self._Tphi
        (rperi,rap)= self.calcRapRperi()
        if rap == rperi:#Circular orbit
            return nu.array([2.*m.pi*self._R/self._vT,0.])
        TR= self.TR(**kwargs)
        I= self.I(**kwargs)
        Tphi= nu.zeros(2)
        Tphi[0]= TR[0]/I[0]*m.pi
        Tphi[1]= Tphi[0]*m.sqrt((I[1]/I[0])**2.+(TR[1]/TR[0])**2.)
        self._Tphi= Tphi
        return self._Tphi

    def I(self,**kwargs):
        """
        NAME:
           I
        PURPOSE:
           Calculate I, the 'ratio' between the radial and azimutha period,
           for a power-law rotation curve
        INPUT:
           +scipy.integrate.quadrature keywords
        OUTPUT:
           I(R,vT,vT) + estimate of the error
        HISTORY:
           2010-11-30 - Written - Bovy (NYU)
        """
        if hasattr(self,'_I'):
            return self._I
        (rperi,rap)= self.calcRapRperi()
        Rmean= m.exp((m.log(rperi)+m.log(rap))/2.)
        if rap == rperi: #Rough limit
            TR= self.TR()[0]
            Tphi= self.Tphi()[0]
            self._I= nu.array([TR/Tphi,0.])
            return self._I
        I= 0.
        if Rmean > rperi:
            I+= rperi**(-1.-self._beta)\
                *nu.array(integrate.quadrature(_IPowerIntegrandSmall,
                                               0.,m.sqrt(Rmean/rperi-1.),
                                               args=(self._X,self._beta,
                                                     self._signbeta),
                                               **kwargs))
        if Rmean < rap:
            I+= rap**(-1.-self._beta)\
                *nu.array(integrate.quadrature(_IPowerIntegrandLarge,
                                               0.,m.sqrt(1.-Rmean/rap),
                                               args=(self._Y,self._beta,
                                                     self._signbeta),
                                               **kwargs))
        self._I= I*self._R*self._vT*m.sqrt(self._absbeta)
        return self._I

    def Jphi(self):
        """
        NAME:
           Jphi
        PURPOSE:
           Calculate the azimuthal action
        INPUT:
        OUTPUT:
           J_R(R,vT,vT)/ro/vc
        HISTORY:
           2010-05-13 - Written - Bovy (NYU)
        """
        return nu.array([self._R*self._vT,0.])

    def JR(self,**kwargs):
        """
        NAME:
           JR
        PURPOSE:
           Calculate the radial action for a power-law rotation curve
        INPUT:
           +scipy.integrate.quad keywords
        OUTPUT:
           J_R(R,vT,vT)/ro/vc + estimate of the error
        HISTORY:
           2010-11-30 - Written - Bovy (NYU)
        """
        if hasattr(self,'_JR'):
            return self._JR
        (rperi,rap)= self.calcRapRperi()
        self._JR= (2./m.sqrt(self._absbeta)*rperi**(self._beta+1.)*
                   nu.array(integrate.quad(_JRPowerIntegrand,1.,rap/rperi,
                                           args=(self._X,self._beta,
                                                 self._signbeta),
                                           **kwargs)))
        return self._JR

    def calcRapRperi(self):
        """
        NAME:
           calcRapRperi
        PURPOSE:
           calculate the apocenter and pericenter radii for a power-law
           rotation curve
        INPUT:
        OUTPUT:
           (rperi,rap)
        HISTORY:
           2010-11-30 - Written - Bovy (NYU)
        """
        if hasattr(self,'_rperirap'):
            return self._rperirap
        EL= calcELPower(self._R,self._vR,self._vT,self._beta,vc=1.,ro=1.)
        E, L= EL
        if self._vR == 0. and self._vT > self._R**self._beta: #We are exactly at pericenter
            rperi= self._R
            rend= _rapRperiPowerFindStart(self._R,E,L,self._beta,rap=True)
            rap= optimize.brentq(_rapRperiPowerEq,rperi+0.000001,
                                 rend,args=(E,L,self._beta))
#                                 fprime=_rapRperiPowerDeriv)
        elif self._vR == 0. and self._vT <self.R**self._beta: #We are exactly at apocenter
            rap= self._R
            rstart= _rapRperiPowerFindStart(self._R,E,L,self._beta)
            rperi= optimize.brentq(_rapRperiPowerEq,rstart,rap-0.000001,
                                   args=(E,L,self._beta))
#                                   fprime=_rapRperiPowerDeriv)
        elif self._vR == 0. and self._vT == self._R**self.__beta: #We are on a circular orbit
            rperi= self._R
            rap = self._R
        else:
            rstart= _rapRperiPowerFindStart(self._R,E,L,self._beta)
            rperi= optimize.brentq(_rapRperiPowerEq,rstart,self._R,
                                   (E,L,self._beta))
            rend= _rapRperiPowerFindStart(self._R,E,L,self._beta,rap=True)
            rap= optimize.brentq(_rapRperiPowerEq,self._R,rend,
                                 (E,L,self._beta))
        self._rperirap= (rperi,rap)
        return self._rperirap

def calcRapRperiFromELPower(E,L,beta,vc=1.,ro=1.):
    """
    NAME:
       calcRapRperiFromELPower
    PURPOSE:
       calculate the apocenter and pericenter radii for a power-law
       rotation curve
    INPUT:
       E - energy
       L - angular momemtum
       beta - power-law index of rotation curve
       vc - circular velocity
       ro - reference radius
    OUTPUT:
       (rperi,rap)
    HISTORY:
       2010-11-30 - Written - Bovy (NYU)
    """
    rstart= _rapRperiPowerFindStart(L,E,L)
    rperi= optimize.brentq(_rapRperiPowerEq,rstart,L,(E,L,beta))
    rend= _rapRperiPowerFindStart(L,E,L,rap=True)
    rap= optimize.brentq(_rapRperiPowerEq,L,rend,(E,L,beta))
    return (rperi,rap)

def calcELPower(R,vR,vT,beta,vc=1.,ro=1.):
    """
    NAME:
       calcELPower
    PURPOSE:
       calculate the energy and angular momentum for a flat rotation curve
    INPUT:
       R - Galactocentric radius (/ro)
       vR - radial part of the velocity (/vc)
       vT - azimuthal part of the velocity (/vc)
       vc - circular velocity
       ro - reference radius
    OUTPUT:
       (E,L)
    HISTORY:
       2010-11-30 - Written - Bovy (NYU)
    """
    return (potentialPower(R,beta,vc=vc,ro=ro)+vR**2./2.+vT**2./2.,R*vT)

def potentialPower(R,beta,vc=1.,ro=1.):
    """
    NAME:
       potentialPower
    PURPOSE:
       return the potential for a power-law rotation curve
    INPUT:
       R - Galactocentric radius (/ro)
       beta - power-law index
       vc - circular velocity
       ro - reference radius
    OUTPUT:
       Phi(R)
    HISTORY:
       2010-11-30 - Written - Bovy (NYU)
    """
    return vc**2./2./beta*(R/ro)**(2.*beta)

def _JRPowerIntegrand(r,X,beta,signbeta):
    """The J_R integrand for a flat rotation curve"""
    return nu.sqrt(signbeta*(1.-X-r**(2.*beta+2.)+r**2.*X))/r

def _TRPowerIntegrandSmall(t,X,beta,signbeta):
    r= 1.+t**2.#part of the transformation
    return 2.*t/_JRPowerIntegrand(r,X,beta,signbeta)

def _TRPowerIntegrandLarge(t,Y,beta,signbeta):
    r= 1.-t**2.#part of the transformation
    return 2.*t/_JRPowerIntegrand(r,Y,beta,signbeta) #same integrand

def _IPowerIntegrandSmall(t,X,beta,signbeta):
    r= 1.+t**2.#part of the transformation
    return 2.*t/_JRPowerIntegrand(r,X,beta,signbeta)/r**2.

def _IPowerIntegrandLarge(t,Y,beta,signbeta):
    r= 1.-t**2.#part of the transformation
    return 2.*t/_JRPowerIntegrand(r,Y,beta,signbeta)/r**2.

def _rapRperiPowerEq(R,E,L,beta):
    """The vr=0 equation that needs to be solved to find apo- and pericenter"""
    return E-potentialPower(R,beta)-L**2./2./R**2.

def _rapRperiPowerDeriv(R,E,L,beta):
    """The derivative of the vr=0 equation that needs to be solved to find
    apo- and pericenter"""
    return -R**(2.*beta-1.)+L**2./R**3.

def _rapRperiPowerFindStart(R,E,L,beta,rap=False):
    """
    NAME:
       _rapRperiPowerFindStart
    PURPOSE:
       Find adequate start or end points to solve for rap and rperi
    INPUT:
       R - Galactocentric radius
       E - energy
       L - angular momentum
       beta - power-law index of rotation curve
       rap - if True, find the rap end-point
    OUTPUT:
       rstart or rend
    HISTORY:
       2010-11-30 - Written - Bovy (NYU)
    """
    if rap:
        rtry= 2.*R
    else:
        rtry= R/2.
    while (E-potentialPower(rtry,beta)-L**2./2./rtry**2) > 0.:
        if rap:
            if rtry > 100.:
                raise UnboundError("Orbit seems to be unboud")
            rtry*= 2.
        else:
            rtry/= 2.
    return rtry

<|endoftext|>


<commit_before># Copyright 2021 Google LLC
#
# 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.
"""GCE nodes have an up to date ops agent installed.

Verify that the ops agent is used by the GCE instances
and that the agent is recent enough.
If the monitoring agent is found it is recommended to upgrade to the ops agent.

see: https://cloud.google.com/stackdriver/docs/solutions/agents/ops-agent
"""

import operator as op
from datetime import datetime, timedelta
from typing import Dict

from packaging import version

from gcpdiag import lint, models
from gcpdiag.queries import gce, monitoring

_query_results_project_id: Dict[str, monitoring.TimeSeriesCollection] = {}

INSTANCE_MIN_AGE = timedelta(hours=1)

OPS_NAME = 'google-cloud-ops-agent-metrics'
MONITORING_NAME = 'stackdriver_agent'
OPS_AGENT_MIN_VERSION = version.parse('2.6.0')

def prefetch_rule(context: models.Context):
  # Fetch the metrics for all instances.
  instances = [
      vm for vm in gce.get_instances(context).values() if not vm.is_gke_node()
  ]
  if not instances:
    return

  _query_results_project_id[context.project_id] = \
      monitoring.query(
          context.project_id, """
fetch gce_instance
| metric 'agent.googleapis.com/agent/uptime'
| align rate(4m)
| every 4m
  """)

def run_rule(context: models.Context, report: lint.LintReportRuleInterface):
  instances = [
      vm for vm in gce.get_instances(context).values() if not vm.is_gke_node()
  ]
  if not instances:
    report.add_skipped(None, 'no instances found')
    return
  if not context.project_id in _query_results_project_id:
    report.add_skipped(None, 'no instances found')
    return
  query = _query_results_project_id[context.project_id]
  try:
    vms_agents = {
        e['labels']['resource.instance_id']: e['labels']['metric.version']
        for e in query.values()
    }
  except KeyError:
    report.add_skipped(None, 'The metrics have incorrect labels')
    return
  for i in sorted(instances, key=op.attrgetter('project_id', 'name')):
    if datetime.utcnow() - i.creation_timestamp < INSTANCE_MIN_AGE:
      #instance is less than 1 hour of creation agent data might not be accurate yet
      continue
    if i.id in vms_agents:
      if MONITORING_NAME in vms_agents[i.id]:
        report.add_failed(
            i,
            'instance using old monitoring agent instead of ops agent. See description to upgrade'
        )
      if OPS_NAME in vms_agents[i.id]:
        if OPS_AGENT_MIN_VERSION <= version.parse(
            vms_agents[i.id].split('/')[1].split('-')[0]):
          report.add_ok(i)
        else:
          # if an unsafe or buggy old version is present
          # we notify the user here by changing to add fail
          report.add_failed(
              i,
              f'ops agent too old, should be at least ({OPS_AGENT_MIN_VERSION})'
          )
    else:
      report.add_failed(i, 'ops agent not installed.')
<commit_msg>gce/BP/002: make reporting more compact<commit_after># Copyright 2021 Google LLC
#
# 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.
"""GCE nodes have an up to date ops agent installed.

Verify that the ops agent is used by the GCE instances
and that the agent is recent enough.
If the monitoring agent is found it is recommended to upgrade to the ops agent.

see: https://cloud.google.com/stackdriver/docs/solutions/agents/ops-agent
"""

import operator as op
from datetime import datetime, timedelta
from typing import Dict

from packaging import version

from gcpdiag import lint, models
from gcpdiag.queries import gce, monitoring

_query_results_project_id: Dict[str, monitoring.TimeSeriesCollection] = {}

INSTANCE_MIN_AGE = timedelta(hours=1)

OPS_NAME = 'google-cloud-ops-agent-metrics'
MONITORING_NAME = 'stackdriver_agent'
OPS_AGENT_MIN_VERSION = version.parse('2.6.0')

def prefetch_rule(context: models.Context):
  # Fetch the metrics for all instances.
  instances = [
      vm for vm in gce.get_instances(context).values() if not vm.is_gke_node()
  ]
  if not instances:
    return

  _query_results_project_id[context.project_id] = \
      monitoring.query(
          context.project_id, """
fetch gce_instance
| metric 'agent.googleapis.com/agent/uptime'
| align rate(4m)
| every 4m
  """)

def run_rule(context: models.Context, report: lint.LintReportRuleInterface):
  instances = [
      vm for vm in gce.get_instances(context).values() if not vm.is_gke_node()
  ]
  if not instances:
    report.add_skipped(None, 'no instances found')
    return
  if not context.project_id in _query_results_project_id:
    report.add_skipped(None, 'no instances found')
    return
  query = _query_results_project_id[context.project_id]
  try:
    vms_agents = {
        e['labels']['resource.instance_id']: e['labels']['metric.version']
        for e in query.values()
    }
  except KeyError:
    report.add_skipped(None, 'The metrics have incorrect labels')
    return
  for i in sorted(instances, key=op.attrgetter('project_id', 'name')):
    if datetime.utcnow() - i.creation_timestamp < INSTANCE_MIN_AGE:
      #instance is less than 1 hour of creation agent data might not be accurate yet
      continue
    if i.id in vms_agents:
      if MONITORING_NAME in vms_agents[i.id]:
        report.add_failed(
            i, 'instance using old monitoring agent instead of ops agent.')
      if OPS_NAME in vms_agents[i.id]:
        if OPS_AGENT_MIN_VERSION <= version.parse(
            vms_agents[i.id].split('/')[1].split('-')[0]):
          report.add_ok(i)
        else:
          # if an unsafe or buggy old version is present
          # we notify the user here by changing to add fail
          version_str = vms_agents[i.id].split('/')[1].split('-')[0]
          report.add_failed(
              i,
              f'ops agent too old: {version_str}, should be at least {OPS_AGENT_MIN_VERSION}'
          )
    else:
      report.add_failed(i, None, 'not installed')
<|endoftext|>


<commit_before>from __future__ import print_function

# define some standard config-free analyzers
standard_analyzers = {
    'simple':       {'tokenizer': 'lowercase',  'filter': []},
    'whitespace':   {'tokenizer': 'whitespace', 'filter': []},
    'keyword':      {'tokenizer': 'keyword',    'filter': []},
    'stop':         {'tokenizer': 'lowercase',  'filter': ['stop']},
    'standard':     {'tokenizer': 'standard',   'filter': ['standard', 'lowercase', 'stop']},
    'snowball':     {'tokenizer': 'standard',   'filter': ['standard', 'lowercase', 'stop', 'snowball']}
}

class AnalyzerNotFound(Exception):
    pass

def normalizeAnalyzer(analyzer):
    try:
        if (isinstance(analyzer['char_filter'], str) or isinstance(analyzer['char_filter'], unicode)):
            analyzer['char_filter'] = [analyzer['char_filter']]
    except KeyError:
        analyzer['char_filter'] = []

def getAnalyzer(indexName, analyzerName, es):
    # try standard analyzers first
    if analyzerName in standard_analyzers:
        analyzer = standard_analyzers[analyzerName]
        normalizeAnalyzer(analyzer)
        return analyzer

    # otherwise try custom ones
    settings = es.indices.get_settings(index=indexName)
    try:
        analyzer = settings[indexName]['settings']['index']['analysis']['analyzer'][analyzerName]
        normalizeAnalyzer(analyzer)
        return analyzer
    except KeyError:
        raise AnalyzerNotFound()

def printTokens(analyzeResp):
    tokens = analyzeResp['tokens']
    posnToTokens = {}
    for token in tokens:
        try:
            posnToTokens[token['position']].append(token['token'])
        except KeyError:
            posnToTokens[token['position']] = [token['token']]

    outputStr = ""
    for position, tokens in posnToTokens.iteritems():
        outputStr += ("{%s:" % position) + ",".join(tokens) + "}\t"
    print(outputStr)

def stepWise(text, indexName, analyzer, es):
    charFiltersInUse = []
    tokenizer = analyzer['tokenizer']
    filtersInUse = []
    # Just each char filter
    for charFilter in analyzer['char_filter']:
        print("CHAR_FILTER: %s" % charFilter)
        charFiltersInUse.append(charFilter)
        body = {"text": text, "char_filter": charFiltersInUse}
        analyzeResp = es.indices.analyze(index=indexName, body=body)
        printTokens(analyzeResp)

    # Add tokenizer
    print("TOKENIZER: %s" % tokenizer)
    print("char filters %s" % charFiltersInUse)
    body = {"text": text, "char_filter": charFiltersInUse, "tokenizer": tokenizer}
    analyzeResp = es.indices.analyze(index=indexName, body=body)
    printTokens(analyzeResp)

    # Token Filters
    filters = []
    if 'filter' in analyzer:
        filters = analyzer['filter']
    elif 'filters' in analyzer:
        filters = analyzer['filters']
    else:
        raise ValueError("Weird... No Filters for analyzer %s" % analyzer)
    print("filters %s" % filters)
    for currFilter in filters:
        print("TOKEN_FILTER: %s" % currFilter)
        filtersInUse.append(currFilter)
        body = {"text": text, "char_filter": charFiltersInUse, "tokenizer": tokenizer, "filter": filtersInUse}
        analyzeResp = es.indices.analyze(index=indexName, body=body)
        printTokens(analyzeResp)
<commit_msg>Removes debug statements<commit_after>from __future__ import print_function

# define some standard config-free analyzers
standard_analyzers = {
    'simple':       {'tokenizer': 'lowercase',  'filter': []},
    'whitespace':   {'tokenizer': 'whitespace', 'filter': []},
    'keyword':      {'tokenizer': 'keyword',    'filter': []},
    'stop':         {'tokenizer': 'lowercase',  'filter': ['stop']},
    'standard':     {'tokenizer': 'standard',   'filter': ['standard', 'lowercase', 'stop']},
    'snowball':     {'tokenizer': 'standard',   'filter': ['standard', 'lowercase', 'stop', 'snowball']}
}

class AnalyzerNotFound(Exception):
    pass

def normalizeAnalyzer(analyzer):
    try:
        if (isinstance(analyzer['char_filter'], str) or isinstance(analyzer['char_filter'], unicode)):
            analyzer['char_filter'] = [analyzer['char_filter']]
    except KeyError:
        analyzer['char_filter'] = []

def getAnalyzer(indexName, analyzerName, es):
    # try standard analyzers first
    if analyzerName in standard_analyzers:
        analyzer = standard_analyzers[analyzerName]
        normalizeAnalyzer(analyzer)
        return analyzer

    # otherwise try custom ones
    settings = es.indices.get_settings(index=indexName)
    try:
        analyzer = settings[indexName]['settings']['index']['analysis']['analyzer'][analyzerName]
        normalizeAnalyzer(analyzer)
        return analyzer
    except KeyError:
        raise AnalyzerNotFound()

def printTokens(analyzeResp):
    tokens = analyzeResp['tokens']
    posnToTokens = {}
    for token in tokens:
        try:
            posnToTokens[token['position']].append(token['token'])
        except KeyError:
            posnToTokens[token['position']] = [token['token']]

    outputStr = ""
    for position, tokens in posnToTokens.iteritems():
        outputStr += ("{%s:" % position) + ",".join(tokens) + "}\t"
    print(outputStr)

def stepWise(text, indexName, analyzer, es):
    charFiltersInUse = []
    tokenizer = analyzer['tokenizer']
    filtersInUse = []
    # Just each char filter
    for charFilter in analyzer['char_filter']:
        print("CHAR_FILTER: %s" % charFilter)
        charFiltersInUse.append(charFilter)
        body = {"text": text, "char_filter": charFiltersInUse}
        analyzeResp = es.indices.analyze(index=indexName, body=body)
        printTokens(analyzeResp)

    # Add tokenizer
    print("TOKENIZER: %s" % tokenizer)
    body = {"text": text, "char_filter": charFiltersInUse, "tokenizer": tokenizer}
    analyzeResp = es.indices.analyze(index=indexName, body=body)
    printTokens(analyzeResp)

    # Token Filters
    filters = []
    if 'filter' in analyzer:
        filters = analyzer['filter']
    elif 'filters' in analyzer:
        filters = analyzer['filters']
    else:
        raise ValueError("Weird... No Filters for analyzer %s" % analyzer)
    for currFilter in filters:
        print("TOKEN_FILTER: %s" % currFilter)
        filtersInUse.append(currFilter)
        body = {"text": text, "char_filter": charFiltersInUse, "tokenizer": tokenizer, "filter": filtersInUse}
        analyzeResp = es.indices.analyze(index=indexName, body=body)
        printTokens(analyzeResp)
<|endoftext|>


<commit_before><commit_msg>Draft: support object without shape in DraftLink<commit_after><|endoftext|>


<commit_before># Copyright (c) 2010-2012, GEM Foundation.
#
# OpenQuake is free software: you can redistribute it and/or modify it
# under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# OpenQuake is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with OpenQuake.  If not, see <http://www.gnu.org/licenses/>.

import getpass
import unittest

import kombu

from nose.plugins.attrib import attr

from openquake.db import models
from openquake.calculators.hazard.event_based import core_next
from openquake.calculators.hazard import general as haz_general
from openquake.utils import stats

from tests.utils import helpers

class EventBasedHazardCalculatorTestCase(unittest.TestCase):
    """
    Tests for the core functionality of the event-based hazard calculator.
    """

    def setUp(self):
        cfg = helpers.get_data_path('event_based_hazard/job.ini')
        self.job = helpers.get_hazard_job(cfg, username=getpass.getuser())
        self.calc = core_next.EventBasedHazardCalculator(self.job)
        models.JobStats.objects.create(oq_job=self.job)

    def test_initialize_ses_db_records(self):
        hc = self.job.hazard_calculation

        # Initialize sources as a setup for the test:
        self.calc.initialize_sources()

        self.calc.initialize_realizations(
            rlz_callbacks=[self.calc.initialize_ses_db_records])

        outputs = models.Output.objects.filter(
            oq_job=self.job, output_type='ses')
        self.assertEqual(2, len(outputs))

        # With this job configuration, we have 2 logic tree realizations.
        lt_rlzs = models.LtRealization.objects.filter(hazard_calculation=hc)
        self.assertEqual(2, len(lt_rlzs))

        for rlz in lt_rlzs:
            sess = models.SES.objects.filter(
                ses_collection__lt_realization=rlz)
            self.assertEqual(hc.ses_per_logic_tree_path, len(sess))

            for ses in sess:
                # The only metadata in in the SES is investigation time.
                self.assertEqual(hc.investigation_time, ses.investigation_time)

    def test_initialize_pr_data_with_ses(self):
        hc = self.job.hazard_calculation

        # Initialize sources as a setup for the test:
        self.calc.initialize_sources()

        self.calc.initialize_realizations(
            rlz_callbacks=[self.calc.initialize_ses_db_records])

        ltr1, ltr2 = models.LtRealization.objects.filter(
            hazard_calculation=hc).order_by("id")

        ltr1.completed_sources = 12
        ltr1.save()

        self.calc.initialize_pr_data()

        total = stats.pk_get(self.calc.job.id, "nhzrd_total")
        self.assertEqual(ltr1.total_sources + ltr2.total_sources, total)
        done = stats.pk_get(self.calc.job.id, "nhzrd_done")
        self.assertEqual(ltr1.completed_sources + ltr2.completed_sources, done)

    def test_initialize_gmf_db_records(self):
        hc = self.job.hazard_calculation

        # Initialize sources as a setup for the test:
        self.calc.initialize_sources()

        self.calc.initialize_realizations(
            rlz_callbacks=[self.calc.initialize_gmf_db_records])

        outputs = models.Output.objects.filter(
            oq_job=self.job, output_type='gmf')
        self.assertEqual(2, len(outputs))

        lt_rlzs = models.LtRealization.objects.filter(hazard_calculation=hc)
        self.assertEqual(2, len(lt_rlzs))

        for rlz in lt_rlzs:
            gmf_sets = models.GmfSet.objects.filter(
                gmf_collection__lt_realization=rlz)
            self.assertEqual(hc.ses_per_logic_tree_path, len(gmf_sets))

            for gmf_set in gmf_sets:
                # The only metadata in a GmfSet is investigation time.
                self.assertEqual(
                    hc.investigation_time, gmf_set.investigation_time)

    def test_initialize_pr_data_with_gmf(self):
        hc = self.job.hazard_calculation

        # Initialize sources as a setup for the test:
        self.calc.initialize_sources()

        self.calc.initialize_realizations(
            rlz_callbacks=[self.calc.initialize_gmf_db_records])

        ltr1, ltr2 = models.LtRealization.objects.filter(
            hazard_calculation=hc).order_by("id")

        ltr1.completed_sources = 13
        ltr1.save()

        self.calc.initialize_pr_data()

        total = stats.pk_get(self.calc.job.id, "nhzrd_total")
        self.assertEqual(ltr1.total_sources + ltr2.total_sources, total)
        done = stats.pk_get(self.calc.job.id, "nhzrd_done")
        self.assertEqual(ltr1.completed_sources + ltr2.completed_sources, done)

    def test_initialize_complete_lt_ses_db_records_branch_enum(self):
        # Set hazard_calculation.number_of_logic_tree_samples = 0
        # This indicates that the `end-branch enumeration` method should be
        # used to carry out the calculation.

        # This test was added primarily for branch coverage (in the case of end
        # branch enum) for the method `initialize_complete_lt_ses_db_records`.
        hc = self.job.hazard_calculation
        hc.number_of_logic_tree_samples = 0

        self.calc.initialize_sources()
        self.calc.initialize_realizations()

        self.calc.initialize_complete_lt_ses_db_records()

        complete_lt_ses = models.SES.objects.get(
            ses_collection__output__oq_job=self.job.id,
            ses_collection__output__output_type='complete_lt_ses',
            complete_logic_tree_ses=True)

        self.assertEqual(250.0, complete_lt_ses.investigation_time)
        self.assertIsNone(complete_lt_ses.ordinal)

    # TODO(LB): This test is becoming a bit epic. Once QA test data is
    # we can probably refactor or replace this test.
    @attr('slow')
    def test_complete_event_based_calculation_cycle(self):
        # * Run `pre_execute()`.
        # * Execute the `stochastic_event_sets` task as a normal function.
        # * Check that the proper results (GMF, SES) were computed.
        # * Finally, call `post_execute()` and verify that `complete logic
        #   tree` artifacts were created.

        # There 4 sources in the test input model; we can test them all with 1
        # task.
        sources_per_task = 4

        self.calc.pre_execute()
        # Test the job stats:
        job_stats = models.JobStats.objects.get(oq_job=self.job.id)
        self.assertEqual(2, job_stats.num_tasks)
        self.assertEqual(121, job_stats.num_sites)
        self.assertEqual(2, job_stats.num_realizations)

        self.job.is_running = True
        self.job.status = 'executing'
        self.job.save()

        hc = self.job.hazard_calculation
        num_sites = len(hc.points_to_compute())

        rlz1, rlz2 = models.LtRealization.objects.filter(
            hazard_calculation=hc.id)

        progress = dict(total=0, computed=0)

        task_arg_gen = self.calc.task_arg_gen(
            hc, self.job, sources_per_task, progress)
        task_arg_list = list(task_arg_gen)

        self.assertEqual(2, len(task_arg_list))

        # Now test the completion signal messaging of the task:
        def test_callback(body, message):
            self.assertEqual(
                dict(job_id=self.job.id, num_sources=sources_per_task), body)
            message.ack()

        exchange, conn_args = haz_general.exchange_and_conn_args()

        routing_key = haz_general.ROUTING_KEY_FMT % dict(job_id=self.job.id)
        task_signal_queue = kombu.Queue(
            'htasks.job.%s' % self.job.id, exchange=exchange,
            routing_key=routing_key, durable=False, auto_delete=True)

        with kombu.BrokerConnection(**conn_args) as conn:
            task_signal_queue(conn.channel()).declare()
            with conn.Consumer(task_signal_queue, callbacks=[test_callback]):
                # call the task as a normal function
                for args in task_arg_list:
                    core_next.ses_and_gmfs(*args)

                    # wait for the completion signal
                    conn.drain_events()

        # Check the 'total' counter (computed by the task arg generator):
        # 2 realizations * 4 sources = 8 total
        self.assertEqual(8, progress['total'])

        # Now check that we saved the right number of ruptures to the DB.
        ruptures1 = models.SESRupture.objects.filter(
            ses__ses_collection__lt_realization=rlz1)
        self.assertEqual(118, ruptures1.count())

        ruptures2 = models.SESRupture.objects.filter(
            ses__ses_collection__lt_realization=rlz2)
        self.assertEqual(92, ruptures2.count())

        # Check that we have the right number of gmf_sets.
        # The correct number is (num_realizations * ses_per_logic_tree_path).
        gmf_sets = models.GmfSet.objects.filter(
            gmf_collection__output__oq_job=self.job.id,
            complete_logic_tree_gmf=False)
        # 2 realizations, 5 ses_per_logic_tree_path
        self.assertEqual(10, gmf_sets.count())

        for imt in hc.intensity_measure_types:
            imt, sa_period, sa_damping = models.parse_imt(imt)
            # Now check that we have the right number of GMFs in the DB.
            for gmf_set in gmf_sets:

                # For each gmf_set, we should have a number of GMF records
                # equal to the numbers of sites in the calculation, _per_ IMT.
                # In this case, that's 121.
                gmfs = models.Gmf.objects.filter(
                    gmf_set=gmf_set, imt=imt, sa_period=sa_period,
                    sa_damping=sa_damping)

                # Sanity check: make sure they all came from the same task:
                task_ord = gmfs[0].task_ordinal
                self.assertTrue(all(x.task_ordinal == task_ord for x in gmfs))

                # Expected number of ruptures:
                exp_n_rups = models.SESRupture.objects.filter(
                    ses__ses_collection__output__oq_job=self.job.id,
                    ses__ordinal=gmf_set.ses_ordinal,
                    task_ordinal=task_ord).count()

                self.assertEqual(121, gmfs.count())
                self.assertTrue(all(len(x.gmvs) == exp_n_rups for x in gmfs))

        # TODO: At some point, we'll need to test the actual values of these
        # ruptures. We'll need to collect QA test data for this.

        # Check the complete logic tree SES and make sure it contains
        # all of the ruptures.
        complete_lt_ses = models.SES.objects.get(
            ses_collection__output__oq_job=self.job.id,
            ses_collection__output__output_type='complete_lt_ses',
            complete_logic_tree_ses=True)

        clt_ses_ruptures = models.SESRupture.objects.filter(
            ses=complete_lt_ses.id)

        self.assertEqual(210, clt_ses_ruptures.count())

        # Test the computed `investigation_time`
        # 2 lt realizations * 5 ses_per_logic_tree_path * 50.0 years
        self.assertEqual(500.0, complete_lt_ses.investigation_time)

        self.assertIsNone(complete_lt_ses.ordinal)

        # Check the complete logic tree GMF set and make sure it contains
        # all of the GMFs.
        complete_lt_gmf = models.GmfSet.objects.get(
            gmf_collection__output__oq_job=self.job.id,
            gmf_collection__output__output_type='complete_lt_gmf',
            complete_logic_tree_gmf=True)

        clt_gmfs = models.GmfNode.objects.filter(
            gmf__gmf_set=complete_lt_gmf.id)
        self.assertEqual(expected_gmfs1 + expected_gmfs2, clt_gmfs.count())

        self.assertEqual(500.0, complete_lt_gmf.investigation_time)
        self.assertIsNone(complete_lt_gmf.ses_ordinal)
<commit_msg>tests/calcs/hazard/event_based/core_next_test:<commit_after># Copyright (c) 2010-2012, GEM Foundation.
#
# OpenQuake is free software: you can redistribute it and/or modify it
# under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# OpenQuake is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with OpenQuake.  If not, see <http://www.gnu.org/licenses/>.

import getpass
import unittest

import kombu

from nose.plugins.attrib import attr

from openquake.db import models
from openquake.calculators.hazard.event_based import core_next
from openquake.calculators.hazard import general as haz_general
from openquake.utils import stats

from tests.utils import helpers

class EventBasedHazardCalculatorTestCase(unittest.TestCase):
    """
    Tests for the core functionality of the event-based hazard calculator.
    """

    def setUp(self):
        cfg = helpers.get_data_path('event_based_hazard/job.ini')
        self.job = helpers.get_hazard_job(cfg, username=getpass.getuser())
        self.calc = core_next.EventBasedHazardCalculator(self.job)
        models.JobStats.objects.create(oq_job=self.job)

    def test_initialize_ses_db_records(self):
        hc = self.job.hazard_calculation

        # Initialize sources as a setup for the test:
        self.calc.initialize_sources()

        self.calc.initialize_realizations(
            rlz_callbacks=[self.calc.initialize_ses_db_records])

        outputs = models.Output.objects.filter(
            oq_job=self.job, output_type='ses')
        self.assertEqual(2, len(outputs))

        # With this job configuration, we have 2 logic tree realizations.
        lt_rlzs = models.LtRealization.objects.filter(hazard_calculation=hc)
        self.assertEqual(2, len(lt_rlzs))

        for rlz in lt_rlzs:
            sess = models.SES.objects.filter(
                ses_collection__lt_realization=rlz)
            self.assertEqual(hc.ses_per_logic_tree_path, len(sess))

            for ses in sess:
                # The only metadata in in the SES is investigation time.
                self.assertEqual(hc.investigation_time, ses.investigation_time)

    def test_initialize_pr_data_with_ses(self):
        hc = self.job.hazard_calculation

        # Initialize sources as a setup for the test:
        self.calc.initialize_sources()

        self.calc.initialize_realizations(
            rlz_callbacks=[self.calc.initialize_ses_db_records])

        ltr1, ltr2 = models.LtRealization.objects.filter(
            hazard_calculation=hc).order_by("id")

        ltr1.completed_sources = 12
        ltr1.save()

        self.calc.initialize_pr_data()

        total = stats.pk_get(self.calc.job.id, "nhzrd_total")
        self.assertEqual(ltr1.total_sources + ltr2.total_sources, total)
        done = stats.pk_get(self.calc.job.id, "nhzrd_done")
        self.assertEqual(ltr1.completed_sources + ltr2.completed_sources, done)

    def test_initialize_gmf_db_records(self):
        hc = self.job.hazard_calculation

        # Initialize sources as a setup for the test:
        self.calc.initialize_sources()

        self.calc.initialize_realizations(
            rlz_callbacks=[self.calc.initialize_gmf_db_records])

        outputs = models.Output.objects.filter(
            oq_job=self.job, output_type='gmf')
        self.assertEqual(2, len(outputs))

        lt_rlzs = models.LtRealization.objects.filter(hazard_calculation=hc)
        self.assertEqual(2, len(lt_rlzs))

        for rlz in lt_rlzs:
            gmf_sets = models.GmfSet.objects.filter(
                gmf_collection__lt_realization=rlz)
            self.assertEqual(hc.ses_per_logic_tree_path, len(gmf_sets))

            for gmf_set in gmf_sets:
                # The only metadata in a GmfSet is investigation time.
                self.assertEqual(
                    hc.investigation_time, gmf_set.investigation_time)

    def test_initialize_pr_data_with_gmf(self):
        hc = self.job.hazard_calculation

        # Initialize sources as a setup for the test:
        self.calc.initialize_sources()

        self.calc.initialize_realizations(
            rlz_callbacks=[self.calc.initialize_gmf_db_records])

        ltr1, ltr2 = models.LtRealization.objects.filter(
            hazard_calculation=hc).order_by("id")

        ltr1.completed_sources = 13
        ltr1.save()

        self.calc.initialize_pr_data()

        total = stats.pk_get(self.calc.job.id, "nhzrd_total")
        self.assertEqual(ltr1.total_sources + ltr2.total_sources, total)
        done = stats.pk_get(self.calc.job.id, "nhzrd_done")
        self.assertEqual(ltr1.completed_sources + ltr2.completed_sources, done)

    def test_initialize_complete_lt_ses_db_records_branch_enum(self):
        # Set hazard_calculation.number_of_logic_tree_samples = 0
        # This indicates that the `end-branch enumeration` method should be
        # used to carry out the calculation.

        # This test was added primarily for branch coverage (in the case of end
        # branch enum) for the method `initialize_complete_lt_ses_db_records`.
        hc = self.job.hazard_calculation
        hc.number_of_logic_tree_samples = 0

        self.calc.initialize_sources()
        self.calc.initialize_realizations()

        self.calc.initialize_complete_lt_ses_db_records()

        complete_lt_ses = models.SES.objects.get(
            ses_collection__output__oq_job=self.job.id,
            ses_collection__output__output_type='complete_lt_ses',
            complete_logic_tree_ses=True)

        self.assertEqual(250.0, complete_lt_ses.investigation_time)
        self.assertIsNone(complete_lt_ses.ordinal)

    # TODO(LB): This test is becoming a bit epic. Once QA test data is
    # we can probably refactor or replace this test.
    @attr('slow')
    def test_complete_event_based_calculation_cycle(self):
        # * Run `pre_execute()`.
        # * Execute the `stochastic_event_sets` task as a normal function.
        # * Check that the proper results (GMF, SES) were computed.
        # * Finally, call `post_execute()` and verify that `complete logic
        #   tree` artifacts were created.

        # There 4 sources in the test input model; we can test them all with 1
        # task.
        sources_per_task = 4

        self.calc.pre_execute()
        # Test the job stats:
        job_stats = models.JobStats.objects.get(oq_job=self.job.id)
        self.assertEqual(2, job_stats.num_tasks)
        self.assertEqual(121, job_stats.num_sites)
        self.assertEqual(2, job_stats.num_realizations)

        self.job.is_running = True
        self.job.status = 'executing'
        self.job.save()

        hc = self.job.hazard_calculation
        num_sites = len(hc.points_to_compute())

        rlz1, rlz2 = models.LtRealization.objects.filter(
            hazard_calculation=hc.id)

        progress = dict(total=0, computed=0)

        task_arg_gen = self.calc.task_arg_gen(
            hc, self.job, sources_per_task, progress)
        task_arg_list = list(task_arg_gen)

        self.assertEqual(2, len(task_arg_list))

        # Now test the completion signal messaging of the task:
        def test_callback(body, message):
            self.assertEqual(
                dict(job_id=self.job.id, num_sources=sources_per_task), body)
            message.ack()

        exchange, conn_args = haz_general.exchange_and_conn_args()

        routing_key = haz_general.ROUTING_KEY_FMT % dict(job_id=self.job.id)
        task_signal_queue = kombu.Queue(
            'htasks.job.%s' % self.job.id, exchange=exchange,
            routing_key=routing_key, durable=False, auto_delete=True)

        with kombu.BrokerConnection(**conn_args) as conn:
            task_signal_queue(conn.channel()).declare()
            with conn.Consumer(task_signal_queue, callbacks=[test_callback]):
                # call the task as a normal function
                for args in task_arg_list:
                    core_next.ses_and_gmfs(*args)

                    # wait for the completion signal
                    conn.drain_events()

        # Check the 'total' counter (computed by the task arg generator):
        # 2 realizations * 4 sources = 8 total
        self.assertEqual(8, progress['total'])

        # Now check that we saved the right number of ruptures to the DB.
        ruptures1 = models.SESRupture.objects.filter(
            ses__ses_collection__lt_realization=rlz1)
        self.assertEqual(118, ruptures1.count())

        ruptures2 = models.SESRupture.objects.filter(
            ses__ses_collection__lt_realization=rlz2)
        self.assertEqual(92, ruptures2.count())

        # Check that we have the right number of gmf_sets.
        # The correct number is (num_realizations * ses_per_logic_tree_path).
        gmf_sets = models.GmfSet.objects.filter(
            gmf_collection__output__oq_job=self.job.id,
            complete_logic_tree_gmf=False)
        # 2 realizations, 5 ses_per_logic_tree_path
        self.assertEqual(10, gmf_sets.count())

        for imt in hc.intensity_measure_types:
            imt, sa_period, sa_damping = models.parse_imt(imt)
            # Now check that we have the right number of GMFs in the DB.
            for gmf_set in gmf_sets:

                # For each gmf_set, we should have a number of GMF records
                # equal to the numbers of sites in the calculation, _per_ IMT.
                # In this case, that's 121.
                gmfs = models.Gmf.objects.filter(
                    gmf_set=gmf_set, imt=imt, sa_period=sa_period,
                    sa_damping=sa_damping)

                # Sanity check: make sure they all came from the same task:
                task_ord = gmfs[0].task_ordinal
                self.assertTrue(all(x.task_ordinal == task_ord for x in gmfs))

                # Expected number of ruptures:
                exp_n_rups = models.SESRupture.objects.filter(
                    ses__ses_collection__output__oq_job=self.job.id,
                    ses__ordinal=gmf_set.ses_ordinal,
                    task_ordinal=task_ord).count()

                self.assertEqual(121, gmfs.count())
                self.assertTrue(all(len(x.gmvs) == exp_n_rups for x in gmfs))

        # TODO: At some point, we'll need to test the actual values of these
        # ruptures. We'll need to collect QA test data for this.

        # Check the complete logic tree SES and make sure it contains
        # all of the ruptures.
        complete_lt_ses = models.SES.objects.get(
            ses_collection__output__oq_job=self.job.id,
            ses_collection__output__output_type='complete_lt_ses',
            complete_logic_tree_ses=True)

        clt_ses_ruptures = models.SESRupture.objects.filter(
            ses=complete_lt_ses.id)

        self.assertEqual(210, clt_ses_ruptures.count())

        # Test the computed `investigation_time`
        # 2 lt realizations * 5 ses_per_logic_tree_path * 50.0 years
        self.assertEqual(500.0, complete_lt_ses.investigation_time)

        self.assertIsNone(complete_lt_ses.ordinal)

        # TODO: Test the complete logic tree GMF
        # The `complete logic tree GMF` collection is computed as an aggregate
        # of all the GMFs for a calculation.
        # Because GMFs take up a lot of space, we don't store a copy of this
        # as we do with SES.
<|endoftext|>


<commit_before># Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you 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.

import unittest
import os
import shutil
from airflow.contrib.operators.kubernetes_pod_operator import KubernetesPodOperator
from airflow.contrib.kubernetes.secret import Secret
from airflow import AirflowException
from kubernetes.client.rest import ApiException
from subprocess import check_call
import json
from airflow.contrib.kubernetes.pod_launcher import PodLauncher
from airflow.contrib.kubernetes.pod import Port
from airflow.contrib.kubernetes.volume_mount import VolumeMount
from airflow.contrib.kubernetes.volume import Volume
from tests.compat import mock

try:
    check_call(["/usr/local/bin/kubectl", "get", "pods"])
except Exception as e:
    if os.environ.get('KUBERNETES_VERSION'):
        raise e
    else:
        raise unittest.SkipTest(
            "Kubernetes integration tests require a minikube cluster;"
            "Skipping tests {}".format(e)
        )

class KubernetesPodOperatorTest(unittest.TestCase):

    @staticmethod
    def test_config_path_move():
        new_config_path = '/tmp/kube_config'
        old_config_path = os.path.expanduser('~/.kube/config')
        shutil.copy(old_config_path, new_config_path)

        k = KubernetesPodOperator(
            namespace='default',
            image="ubuntu:16.04",
            cmds=["bash", "-cx"],
            arguments=["echo 10"],
            labels={"foo": "bar"},
            name="test",
            task_id="task",
            config_file=new_config_path
        )
        k.execute(None)

    @mock.patch("airflow.contrib.kubernetes.pod_launcher.PodLauncher.run_pod")
    @mock.patch("airflow.contrib.kubernetes.kube_client.get_kube_client")
    def test_config_path(self, client_mock, launcher_mock):
        from airflow.utils.state import State

        file_path = "/tmp/fake_file"
        k = KubernetesPodOperator(
            namespace='default',
            image="ubuntu:16.04",
            cmds=["bash", "-cx"],
            arguments=["echo 10"],
            labels={"foo": "bar"},
            name="test",
            task_id="task",
            config_file=file_path,
            in_cluster=False,
            cluster_context='default'
        )
        launcher_mock.return_value = (State.SUCCESS, None)
        k.execute(None)
        client_mock.assert_called_with(in_cluster=False,
                                       cluster_context='default',
                                       config_file=file_path)

    @mock.patch("airflow.contrib.kubernetes.pod_launcher.PodLauncher.run_pod")
    @mock.patch("airflow.contrib.kubernetes.kube_client.get_kube_client")
    def test_image_pull_secrets_correctly_set(self, client_mock, launcher_mock):
        from airflow.utils.state import State

        fake_pull_secrets = "fakeSecret"
        k = KubernetesPodOperator(
            namespace='default',
            image="ubuntu:16.04",
            cmds=["bash", "-cx"],
            arguments=["echo 10"],
            labels={"foo": "bar"},
            name="test",
            task_id="task",
            image_pull_secrets=fake_pull_secrets,
            in_cluster=False,
            cluster_context='default'
        )
        launcher_mock.return_value = (State.SUCCESS, None)
        k.execute(None)
        self.assertEqual(launcher_mock.call_args[0][0].image_pull_secrets,
                         fake_pull_secrets)

    @mock.patch("airflow.contrib.kubernetes.pod_launcher.PodLauncher.run_pod")
    @mock.patch("airflow.contrib.kubernetes.pod_launcher.PodLauncher.delete_pod")
    @mock.patch("airflow.contrib.kubernetes.kube_client.get_kube_client")
    def test_pod_delete_even_on_launcher_error(self, client_mock, delete_pod_mock, run_pod_mock):
        k = KubernetesPodOperator(
            namespace='default',
            image="ubuntu:16.04",
            cmds=["bash", "-cx"],
            arguments=["echo 10"],
            labels={"foo": "bar"},
            name="test",
            task_id="task",
            in_cluster=False,
            cluster_context='default',
            is_delete_operator_pod=True
        )
        run_pod_mock.side_effect = AirflowException('fake failure')
        with self.assertRaises(AirflowException):
            k.execute(None)
        assert delete_pod_mock.called

    @staticmethod
    def test_working_pod():
        k = KubernetesPodOperator(
            namespace='default',
            image="ubuntu:16.04",
            cmds=["bash", "-cx"],
            arguments=["echo 10"],
            labels={"foo": "bar"},
            name="test",
            task_id="task"
        )
        k.execute(None)

    @staticmethod
    def test_delete_operator_pod():
        k = KubernetesPodOperator(
            namespace='default',
            image="ubuntu:16.04",
            cmds=["bash", "-cx"],
            arguments=["echo 10"],
            labels={"foo": "bar"},
            name="test",
            task_id="task",
            is_delete_operator_pod=True
        )
        k.execute(None)

    @staticmethod
    def test_pod_hostnetwork():
        k = KubernetesPodOperator(
            namespace='default',
            image="ubuntu:16.04",
            cmds=["bash", "-cx"],
            arguments=["echo 10"],
            labels={"foo": "bar"},
            name="test",
            task_id="task",
            hostnetwork=True
        )
        k.execute(None)

    @staticmethod
    def test_pod_dnspolicy():
        k = KubernetesPodOperator(
            namespace='default',
            image="ubuntu:16.04",
            cmds=["bash", "-cx"],
            arguments=["echo 10"],
            labels={"foo": "bar"},
            name="test",
            task_id="task",
            hostnetwork=True,
            dnspolicy="ClusterFirstWithHostNet"
        )
        k.execute(None)

    @staticmethod
    def test_pod_node_selectors():
        node_selectors = {
            'beta.kubernetes.io/os': 'linux'
        }
        k = KubernetesPodOperator(
            namespace='default',
            image="ubuntu:16.04",
            cmds=["bash", "-cx"],
            arguments=["echo", "10"],
            labels={"foo": "bar"},
            name="test",
            task_id="task",
            node_selectors=node_selectors,
            executor_config={'KubernetesExecutor': {'node_selectors': node_selectors}}
        )
        k.execute(None)

    @staticmethod
    def test_pod_resources():
        resources = {
            'limit_memory': 1,
            'limit_cpu': 1,
        }
        k = KubernetesPodOperator(
            namespace='default',
            image="ubuntu:16.04",
            cmds=["bash", "-cx"],
            arguments=["echo", "10"],
            labels={"foo": "bar"},
            name="test",
            task_id="task",
            resources=resources,
            executor_config={'KubernetesExecutor': {'resources': resources}}
        )
        k.execute(None)

    @staticmethod
    def test_pod_affinity():
        affinity = {
            'nodeAffinity': {
                'requiredDuringSchedulingIgnoredDuringExecution': {
                    'nodeSelectorTerms': [
                        {
                            'matchExpressions': [
                                {
                                    'key': 'beta.kubernetes.io/os',
                                    'operator': 'In',
                                    'values': ['linux']
                                }
                            ]
                        }
                    ]
                }
            }
        }
        k = KubernetesPodOperator(
            namespace='default',
            image="ubuntu:16.04",
            cmds=["bash", "-cx"],
            arguments=["echo", "10"],
            labels={"foo": "bar"},
            name="test",
            task_id="task",
            affinity=affinity,
            executor_config={'KubernetesExecutor': {'affinity': affinity}}
        )
        k.execute(None)

    @staticmethod
    def test_logging():
        with mock.patch.object(PodLauncher, 'log') as mock_logger:
            k = KubernetesPodOperator(
                namespace='default',
                image="ubuntu:16.04",
                cmds=["bash", "-cx"],
                arguments=["echo 10"],
                labels={"foo": "bar"},
                name="test",
                task_id="task",
                get_logs=True
            )
            k.execute(None)
            mock_logger.info.assert_any_call(b"+ echo 10\n")

    @staticmethod
    def test_port():
        port = Port('http', 80)

        k = KubernetesPodOperator(
            namespace='default',
            image="ubuntu:16.04",
            cmds=["bash", "-cx"],
            arguments=["echo 10"],
            labels={"foo": "bar"},
            name="test",
            task_id="task",
            ports=[port]
        )
        k.execute(None)

    @staticmethod
    def test_volume_mount():
        with mock.patch.object(PodLauncher, 'log') as mock_logger:
            volume_mount = VolumeMount('test-volume',
                                       mount_path='/root/mount_file',
                                       sub_path=None,
                                       read_only=True)

            volume_config = {
                'persistentVolumeClaim':
                    {
                        'claimName': 'test-volume'
                    }
            }
            volume = Volume(name='test-volume', configs=volume_config)
            k = KubernetesPodOperator(
                namespace='default',
                image="ubuntu:16.04",
                cmds=["bash", "-cx"],
                arguments=["cat /root/mount_file/test.txt"],
                labels={"foo": "bar"},
                volume_mounts=[volume_mount],
                volumes=[volume],
                name="test",
                task_id="task"
            )
            k.execute(None)
            mock_logger.info.assert_any_call(b"retrieved from mount\n")

    @staticmethod
    def test_run_as_user_root():
        security_context = {
            'securityContext': {
                'runAsUser': 0,
            }
        }

        k = KubernetesPodOperator(
            namespace='default',
            image="ubuntu:16.04",
            cmds=["bash", "-cx"],
            arguments=["echo", "10"],
            labels={"foo": "bar"},
            name="test",
            task_id="task",
            security_context=security_context,
            executor_config={'KubernetesExecutor': {'securityContext': security_context}}
        )
        k.execute(None)

    @staticmethod
    def test_run_as_user_non_root():
        security_context = {
            'securityContext': {
                'runAsUser': 1000,
            }
        }

        k = KubernetesPodOperator(
            namespace='default',
            image="ubuntu:16.04",
            cmds=["bash", "-cx"],
            arguments=["echo", "10"],
            labels={"foo": "bar"},
            name="test",
            task_id="task",
            security_context=security_context,
            executor_config={'KubernetesExecutor': {'securityContext': security_context}}
        )
        k.execute(None)

    @staticmethod
    def test_fs_group():
        security_context = {
            'securityContext': {
                'fsGroup': 1000,
            }
        }

        k = KubernetesPodOperator(
            namespace='default',
            image="ubuntu:16.04",
            cmds=["bash", "-cx"],
            arguments=["echo", "10"],
            labels={"foo": "bar"},
            name="test",
            task_id="task",
            security_context=security_context,
            executor_config={'KubernetesExecutor': {'securityContext': security_context}}
        )
        k.execute(None)

    def test_faulty_image(self):
        bad_image_name = "foobar"
        k = KubernetesPodOperator(
            namespace='default',
            image=bad_image_name,
            cmds=["bash", "-cx"],
            arguments=["echo 10"],
            labels={"foo": "bar"},
            name="test",
            task_id="task",
            startup_timeout_seconds=5
        )
        with self.assertRaises(AirflowException):
            k.execute(None)

    def test_faulty_service_account(self):
        bad_service_account_name = "foobar"
        k = KubernetesPodOperator(
            namespace='default',
            image="ubuntu:16.04",
            cmds=["bash", "-cx"],
            arguments=["echo 10"],
            labels={"foo": "bar"},
            name="test",
            task_id="task",
            startup_timeout_seconds=5,
            service_account_name=bad_service_account_name
        )
        with self.assertRaises(ApiException):
            k.execute(None)

    def test_pod_failure(self):
        """
            Tests that the task fails when a pod reports a failure
        """
        bad_internal_command = "foobar"
        k = KubernetesPodOperator(
            namespace='default',
            image="ubuntu:16.04",
            cmds=["bash", "-cx"],
            arguments=[bad_internal_command + " 10"],
            labels={"foo": "bar"},
            name="test",
            task_id="task"
        )
        with self.assertRaises(AirflowException):
            k.execute(None)

    def test_xcom_push(self):
        return_value = '{"foo": "bar"\n, "buzz": 2}'
        k = KubernetesPodOperator(
            namespace='default',
            image="ubuntu:16.04",
            cmds=["bash", "-cx"],
            arguments=['echo \'{}\' > /airflow/xcom/return.json'.format(return_value)],
            labels={"foo": "bar"},
            name="test",
            task_id="task",
            xcom_push=True
        )
        self.assertEqual(k.execute(None), json.loads(return_value))

    @mock.patch("airflow.contrib.kubernetes.pod_launcher.PodLauncher.run_pod")
    @mock.patch("airflow.contrib.kubernetes.kube_client.get_kube_client")
    def test_envs_from_configmaps(self, client_mock, launcher_mock):
        # GIVEN
        from airflow.utils.state import State
        configmaps = ['test-configmap']
        # WHEN
        k = KubernetesPodOperator(
            namespace='default',
            image="ubuntu:16.04",
            cmds=["bash", "-cx"],
            arguments=["echo 10"],
            labels={"foo": "bar"},
            name="test",
            task_id="task",
            configmaps=configmaps
        )
        # THEN
        launcher_mock.return_value = (State.SUCCESS, None)
        k.execute(None)
        self.assertEqual(launcher_mock.call_args[0][0].configmaps, configmaps)

    @mock.patch("airflow.contrib.kubernetes.pod_launcher.PodLauncher.run_pod")
    @mock.patch("airflow.contrib.kubernetes.kube_client.get_kube_client")
    def test_envs_from_secrets(self, client_mock, launcher_mock):
        # GIVEN
        from airflow.utils.state import State
        secrets = [Secret('env', None, "secret_name")]
        # WHEN
        k = KubernetesPodOperator(
            namespace='default',
            image="ubuntu:16.04",
            cmds=["bash", "-cx"],
            arguments=["echo 10"],
            secrets=secrets,
            labels={"foo": "bar"},
            name="test",
            task_id="task",
        )
        # THEN
        launcher_mock.return_value = (State.SUCCESS, None)
        k.execute(None)
        self.assertEqual(launcher_mock.call_args[0][0].secrets, secrets)

if __name__ == '__main__':
    unittest.main()
<commit_msg>[AIRFLOW-XXX] fix copy/pasta in k8s request factory extract resources (#5657)<commit_after># Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you 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.

import unittest
import os
import shutil
from airflow.contrib.operators.kubernetes_pod_operator import KubernetesPodOperator
from airflow.contrib.kubernetes.secret import Secret
from airflow import AirflowException
from kubernetes.client.rest import ApiException
from subprocess import check_call
import json
from airflow.contrib.kubernetes.pod_launcher import PodLauncher
from airflow.contrib.kubernetes.pod import Port
from airflow.contrib.kubernetes.volume_mount import VolumeMount
from airflow.contrib.kubernetes.volume import Volume
from tests.compat import mock

try:
    check_call(["/usr/local/bin/kubectl", "get", "pods"])
except Exception as e:
    if os.environ.get('KUBERNETES_VERSION'):
        raise e
    else:
        raise unittest.SkipTest(
            "Kubernetes integration tests require a minikube cluster;"
            "Skipping tests {}".format(e)
        )

class KubernetesPodOperatorTest(unittest.TestCase):

    @staticmethod
    def test_config_path_move():
        new_config_path = '/tmp/kube_config'
        old_config_path = os.path.expanduser('~/.kube/config')
        shutil.copy(old_config_path, new_config_path)

        k = KubernetesPodOperator(
            namespace='default',
            image="ubuntu:16.04",
            cmds=["bash", "-cx"],
            arguments=["echo 10"],
            labels={"foo": "bar"},
            name="test",
            task_id="task",
            config_file=new_config_path
        )
        k.execute(None)

    @mock.patch("airflow.contrib.kubernetes.pod_launcher.PodLauncher.run_pod")
    @mock.patch("airflow.contrib.kubernetes.kube_client.get_kube_client")
    def test_config_path(self, client_mock, launcher_mock):
        from airflow.utils.state import State

        file_path = "/tmp/fake_file"
        k = KubernetesPodOperator(
            namespace='default',
            image="ubuntu:16.04",
            cmds=["bash", "-cx"],
            arguments=["echo 10"],
            labels={"foo": "bar"},
            name="test",
            task_id="task",
            config_file=file_path,
            in_cluster=False,
            cluster_context='default'
        )
        launcher_mock.return_value = (State.SUCCESS, None)
        k.execute(None)
        client_mock.assert_called_with(in_cluster=False,
                                       cluster_context='default',
                                       config_file=file_path)

    @mock.patch("airflow.contrib.kubernetes.pod_launcher.PodLauncher.run_pod")
    @mock.patch("airflow.contrib.kubernetes.kube_client.get_kube_client")
    def test_image_pull_secrets_correctly_set(self, client_mock, launcher_mock):
        from airflow.utils.state import State

        fake_pull_secrets = "fakeSecret"
        k = KubernetesPodOperator(
            namespace='default',
            image="ubuntu:16.04",
            cmds=["bash", "-cx"],
            arguments=["echo 10"],
            labels={"foo": "bar"},
            name="test",
            task_id="task",
            image_pull_secrets=fake_pull_secrets,
            in_cluster=False,
            cluster_context='default'
        )
        launcher_mock.return_value = (State.SUCCESS, None)
        k.execute(None)
        self.assertEqual(launcher_mock.call_args[0][0].image_pull_secrets,
                         fake_pull_secrets)

    @mock.patch("airflow.contrib.kubernetes.pod_launcher.PodLauncher.run_pod")
    @mock.patch("airflow.contrib.kubernetes.pod_launcher.PodLauncher.delete_pod")
    @mock.patch("airflow.contrib.kubernetes.kube_client.get_kube_client")
    def test_pod_delete_even_on_launcher_error(self, client_mock, delete_pod_mock, run_pod_mock):
        k = KubernetesPodOperator(
            namespace='default',
            image="ubuntu:16.04",
            cmds=["bash", "-cx"],
            arguments=["echo 10"],
            labels={"foo": "bar"},
            name="test",
            task_id="task",
            in_cluster=False,
            cluster_context='default',
            is_delete_operator_pod=True
        )
        run_pod_mock.side_effect = AirflowException('fake failure')
        with self.assertRaises(AirflowException):
            k.execute(None)
        assert delete_pod_mock.called

    @staticmethod
    def test_working_pod():
        k = KubernetesPodOperator(
            namespace='default',
            image="ubuntu:16.04",
            cmds=["bash", "-cx"],
            arguments=["echo 10"],
            labels={"foo": "bar"},
            name="test",
            task_id="task"
        )
        k.execute(None)

    @staticmethod
    def test_delete_operator_pod():
        k = KubernetesPodOperator(
            namespace='default',
            image="ubuntu:16.04",
            cmds=["bash", "-cx"],
            arguments=["echo 10"],
            labels={"foo": "bar"},
            name="test",
            task_id="task",
            is_delete_operator_pod=True
        )
        k.execute(None)

    @staticmethod
    def test_pod_hostnetwork():
        k = KubernetesPodOperator(
            namespace='default',
            image="ubuntu:16.04",
            cmds=["bash", "-cx"],
            arguments=["echo 10"],
            labels={"foo": "bar"},
            name="test",
            task_id="task",
            hostnetwork=True
        )
        k.execute(None)

    @staticmethod
    def test_pod_dnspolicy():
        k = KubernetesPodOperator(
            namespace='default',
            image="ubuntu:16.04",
            cmds=["bash", "-cx"],
            arguments=["echo 10"],
            labels={"foo": "bar"},
            name="test",
            task_id="task",
            hostnetwork=True,
            dnspolicy="ClusterFirstWithHostNet"
        )
        k.execute(None)

    @staticmethod
    def test_pod_node_selectors():
        node_selectors = {
            'beta.kubernetes.io/os': 'linux'
        }
        k = KubernetesPodOperator(
            namespace='default',
            image="ubuntu:16.04",
            cmds=["bash", "-cx"],
            arguments=["echo", "10"],
            labels={"foo": "bar"},
            name="test",
            task_id="task",
            node_selectors=node_selectors,
            executor_config={'KubernetesExecutor': {'node_selectors': node_selectors}}
        )
        k.execute(None)

    @staticmethod
    def test_pod_resources():
        resources = {}
        k = KubernetesPodOperator(
            namespace='default',
            image="ubuntu:16.04",
            cmds=["bash", "-cx"],
            arguments=["echo", "10"],
            labels={"foo": "bar"},
            name="test",
            task_id="task",
            resources=resources,
            executor_config={'KubernetesExecutor': {'resources': resources}}
        )
        k.execute(None)

    @staticmethod
    def test_pod_affinity():
        affinity = {
            'nodeAffinity': {
                'requiredDuringSchedulingIgnoredDuringExecution': {
                    'nodeSelectorTerms': [
                        {
                            'matchExpressions': [
                                {
                                    'key': 'beta.kubernetes.io/os',
                                    'operator': 'In',
                                    'values': ['linux']
                                }
                            ]
                        }
                    ]
                }
            }
        }
        k = KubernetesPodOperator(
            namespace='default',
            image="ubuntu:16.04",
            cmds=["bash", "-cx"],
            arguments=["echo", "10"],
            labels={"foo": "bar"},
            name="test",
            task_id="task",
            affinity=affinity,
            executor_config={'KubernetesExecutor': {'affinity': affinity}}
        )
        k.execute(None)

    @staticmethod
    def test_logging():
        with mock.patch.object(PodLauncher, 'log') as mock_logger:
            k = KubernetesPodOperator(
                namespace='default',
                image="ubuntu:16.04",
                cmds=["bash", "-cx"],
                arguments=["echo 10"],
                labels={"foo": "bar"},
                name="test",
                task_id="task",
                get_logs=True
            )
            k.execute(None)
            mock_logger.info.assert_any_call(b"+ echo 10\n")

    @staticmethod
    def test_port():
        port = Port('http', 80)

        k = KubernetesPodOperator(
            namespace='default',
            image="ubuntu:16.04",
            cmds=["bash", "-cx"],
            arguments=["echo 10"],
            labels={"foo": "bar"},
            name="test",
            task_id="task",
            ports=[port]
        )
        k.execute(None)

    @staticmethod
    def test_volume_mount():
        with mock.patch.object(PodLauncher, 'log') as mock_logger:
            volume_mount = VolumeMount('test-volume',
                                       mount_path='/root/mount_file',
                                       sub_path=None,
                                       read_only=True)

            volume_config = {
                'persistentVolumeClaim':
                    {
                        'claimName': 'test-volume'
                    }
            }
            volume = Volume(name='test-volume', configs=volume_config)
            k = KubernetesPodOperator(
                namespace='default',
                image="ubuntu:16.04",
                cmds=["bash", "-cx"],
                arguments=["cat /root/mount_file/test.txt"],
                labels={"foo": "bar"},
                volume_mounts=[volume_mount],
                volumes=[volume],
                name="test",
                task_id="task"
            )
            k.execute(None)
            mock_logger.info.assert_any_call(b"retrieved from mount\n")

    @staticmethod
    def test_run_as_user_root():
        security_context = {
            'securityContext': {
                'runAsUser': 0,
            }
        }

        k = KubernetesPodOperator(
            namespace='default',
            image="ubuntu:16.04",
            cmds=["bash", "-cx"],
            arguments=["echo", "10"],
            labels={"foo": "bar"},
            name="test",
            task_id="task",
            security_context=security_context,
            executor_config={'KubernetesExecutor': {'securityContext': security_context}}
        )
        k.execute(None)

    @staticmethod
    def test_run_as_user_non_root():
        security_context = {
            'securityContext': {
                'runAsUser': 1000,
            }
        }

        k = KubernetesPodOperator(
            namespace='default',
            image="ubuntu:16.04",
            cmds=["bash", "-cx"],
            arguments=["echo", "10"],
            labels={"foo": "bar"},
            name="test",
            task_id="task",
            security_context=security_context,
            executor_config={'KubernetesExecutor': {'securityContext': security_context}}
        )
        k.execute(None)

    @staticmethod
    def test_fs_group():
        security_context = {
            'securityContext': {
                'fsGroup': 1000,
            }
        }

        k = KubernetesPodOperator(
            namespace='default',
            image="ubuntu:16.04",
            cmds=["bash", "-cx"],
            arguments=["echo", "10"],
            labels={"foo": "bar"},
            name="test",
            task_id="task",
            security_context=security_context,
            executor_config={'KubernetesExecutor': {'securityContext': security_context}}
        )
        k.execute(None)

    def test_faulty_image(self):
        bad_image_name = "foobar"
        k = KubernetesPodOperator(
            namespace='default',
            image=bad_image_name,
            cmds=["bash", "-cx"],
            arguments=["echo 10"],
            labels={"foo": "bar"},
            name="test",
            task_id="task",
            startup_timeout_seconds=5
        )
        with self.assertRaises(AirflowException):
            k.execute(None)

    def test_faulty_service_account(self):
        bad_service_account_name = "foobar"
        k = KubernetesPodOperator(
            namespace='default',
            image="ubuntu:16.04",
            cmds=["bash", "-cx"],
            arguments=["echo 10"],
            labels={"foo": "bar"},
            name="test",
            task_id="task",
            startup_timeout_seconds=5,
            service_account_name=bad_service_account_name
        )
        with self.assertRaises(ApiException):
            k.execute(None)

    def test_pod_failure(self):
        """
            Tests that the task fails when a pod reports a failure
        """
        bad_internal_command = "foobar"
        k = KubernetesPodOperator(
            namespace='default',
            image="ubuntu:16.04",
            cmds=["bash", "-cx"],
            arguments=[bad_internal_command + " 10"],
            labels={"foo": "bar"},
            name="test",
            task_id="task"
        )
        with self.assertRaises(AirflowException):
            k.execute(None)

    def test_xcom_push(self):
        return_value = '{"foo": "bar"\n, "buzz": 2}'
        k = KubernetesPodOperator(
            namespace='default',
            image="ubuntu:16.04",
            cmds=["bash", "-cx"],
            arguments=['echo \'{}\' > /airflow/xcom/return.json'.format(return_value)],
            labels={"foo": "bar"},
            name="test",
            task_id="task",
            xcom_push=True
        )
        self.assertEqual(k.execute(None), json.loads(return_value))

    @mock.patch("airflow.contrib.kubernetes.pod_launcher.PodLauncher.run_pod")
    @mock.patch("airflow.contrib.kubernetes.kube_client.get_kube_client")
    def test_envs_from_configmaps(self, client_mock, launcher_mock):
        # GIVEN
        from airflow.utils.state import State
        configmaps = ['test-configmap']
        # WHEN
        k = KubernetesPodOperator(
            namespace='default',
            image="ubuntu:16.04",
            cmds=["bash", "-cx"],
            arguments=["echo 10"],
            labels={"foo": "bar"},
            name="test",
            task_id="task",
            configmaps=configmaps
        )
        # THEN
        launcher_mock.return_value = (State.SUCCESS, None)
        k.execute(None)
        self.assertEqual(launcher_mock.call_args[0][0].configmaps, configmaps)

    @mock.patch("airflow.contrib.kubernetes.pod_launcher.PodLauncher.run_pod")
    @mock.patch("airflow.contrib.kubernetes.kube_client.get_kube_client")
    def test_envs_from_secrets(self, client_mock, launcher_mock):
        # GIVEN
        from airflow.utils.state import State
        secrets = [Secret('env', None, "secret_name")]
        # WHEN
        k = KubernetesPodOperator(
            namespace='default',
            image="ubuntu:16.04",
            cmds=["bash", "-cx"],
            arguments=["echo 10"],
            secrets=secrets,
            labels={"foo": "bar"},
            name="test",
            task_id="task",
        )
        # THEN
        launcher_mock.return_value = (State.SUCCESS, None)
        k.execute(None)
        self.assertEqual(launcher_mock.call_args[0][0].secrets, secrets)

if __name__ == '__main__':
    unittest.main()
<|endoftext|>