Home > database >  Pytest: assert a list contains n instances of a class / object
Pytest: assert a list contains n instances of a class / object

Time:01-04

@dataclass
class Component:
    name: str

class Assembly:
    def __init__():
        names = ["Able", "Baker", "Charlie"]
        self.components = [Component(name) for name in names]

def test_assembly():
    """Assembly should be initialised with three Components"""
    assembly = Assembly()
    assert assembly.components == [Component] * 3 # doesn't work

The assert doesn't work because it's comparing a list of class instances (assembly.components) to a list of class types. Is there a good way to write this test?

CodePudding user response:

You can convert all of the elements in assembly.components into types with map() and then you can convert the map into a list and use the list.count function to see how many times the type appears in the list. Something like this:

assert list(map(type,assembly.components)).count(Component) == 3

if each item in assembly.components is already a type, you can just make sure its a list and then use the list.count function immedietly like this:

assert list(assembly.components).count(Component) == 3

Edit: Also should you find yourself iterating over anything and want to compare the number of times a pattern appear, i believe its far better to use an int variable and increment it for each pattern you find rather than what you were doing which is creating entirely new lists and comparing them which in large quantities will slow down quite quick. (The increment method is how i bleieve the .count function works so where possible try to just use that.)

CodePudding user response:

You could iterate through the list and use isinstance():

def count_instance(lst, cls):
    count = 0
    for i in lst:
        if isinstance(i, cls):
            count  = 1
    return count
assert count_instance(assembly.components, Component) == 3

CodePudding user response:

assert [type(comp) for comp in assembly.components] == [Component] * 3

This does the job - but I suspect there may be a better solution?

  •  Tags:  
  • Related