Search code examples
pythonpython-3.xunit-testingpython-dataclasses

How to use a Python Dataclass in another class


I'm trying to get to grips with Python and seem to be hitting a wall when trying to use Dataclasses. But when I run the test I have for it I get assertion error as it doesn't seem to see the dataclass right.

I have the following code:

file: music_library.py

from dataclasses import dataclass

@dataclass
class Track:
    title: str
    artist: str
    file: str

class MusicLibrary:
    def __init__(self):
        self.track = Track

    def all(self):
        return self.track

    def add(self, title, artist, file):
        self.track(title = title, artist = artist, file = file)

the add function is being called from a test and is passed the three arguments:

import unittest

from player.music_library import MusicLibrary


class TestMusicLibrary(unittest.TestCase):

    ml = MusicLibrary()

    def test_all(self):
        ml = MusicLibrary()
        ml.add("Track1", "artist1","file1")
        self.assertEqual(ml.all(), ["Track1","artist1","file1" ])

Yet the test fails with

Traceback (most recent call last):
  File "/projects/python/python-music-player-challenges/seed/tests/test_music_library.py", line 13, in test_all
    self.assertEqual(ml.all(), ["Track1","artist1","file1" ])
AssertionError: <class 'player.music_library.Track'> != ['Track1', 'artist1', 'file1']

What's going on here? I'm obviously missing something obvious.

Thanks


Solution

  • Update music_library.py like this:

    from dataclasses import dataclass
    
    @dataclass
    class Track:
        title: str
        artist: str
        file: str
    
    
    class MusicLibrary:
        def __init__(self):
            self.track = None
    
        def all(self):
            return self.track
    
        def add(self, title, artist, file):
            self.track = Track(title=title, artist=artist, file=file)
    
    

    Pay attention to Dataclass instantiation in the above code.

    And update your test case like this:

    import unittest
    
    from music_library import MusicLibrary
    
    
    class TestMusicLibrary(unittest.TestCase):
    
        def test_all(self):
            ml = MusicLibrary()
            ml.add("Track1", "artist1", "file1")
            self.assertEqual([ml.all().title, ml.all().artist, ml.all().file],
                             ["Track1", "artist1", "file1"])
    

    In your test code, you are comparing different object types, you should convert output of ml.all() to list first.

    And if you run the test you get the following output:

    Ran 1 test in 0.000s
    
    OK